feat(config): register diffusion_models as a first-class path key
ComfyUI keeps UNet-only Flux checkpoints under models/diffusion_models/ (paired with separate clip/ and vae/ files), distinct from the unified models/checkpoints/. tsr previously had no way to track that directory in its config — `tsr config --set-path diffusion_models=...` errored because the CLI validator's allow-list only knew the 7 CivitAI-aligned types (checkpoints, loras, embeddings, vae, controlnet, upscalers, other). Workaround was a manual symlink from diffusion_models/ back into checkpoints/, which is fragile (mixed link styles, breaks if the target file is reorganised) and confuses `tsr models`/`tsr db list` (same file shows up under two loader headings). Changes: - config.py: new VALID_PATH_TYPES constant as the single source of truth for keys accepted by `tsr config --set-path KEY=PATH`. Adds "diffusion_models" to the list and a synthetic "DiffusionModel" bucket in DEFAULT_PATHS pointing at MODELS_DIR/diffusion_models. Comment notes that CivitAI does not have a DiffusionModel model_type (UNet-only Flux files are returned as "Checkpoint"), so auto-routing from `tsr dl` will never pick this bucket — it exists for explicit `tsr dl -o /path/to/diffusion_models` workflows and for `tsr config --show` visibility into where the dir lives on disk. - cli.py: imports VALID_PATH_TYPES and replaces the two duplicated inline lists (set-path validator + custom-marker display in the `config` command). No behaviour change for the existing 7 types. After this lands and is `uv tool upgrade tensors`'d, users can `tsr config --set-path diffusion_models=/path/to/comfyui/models/diffusion_models` and `tsr dl -v <civitai-version-id> -o /path/to/diffusion_models` for UNet-only Flux downloads, no more symlink workarounds.
This commit is contained in:
+4
-7
@@ -29,6 +29,7 @@ from tensors.config import (
|
||||
COMFYUI_DEFAULT_WIDTH,
|
||||
CONFIG_FILE,
|
||||
MODEL_FAMILY_DEFAULTS,
|
||||
VALID_PATH_TYPES,
|
||||
BaseModel,
|
||||
CommercialUse,
|
||||
ModelType,
|
||||
@@ -673,10 +674,9 @@ def config(
|
||||
|
||||
path_type, path_value = set_path.split("=", 1)
|
||||
path_type = path_type.lower().strip()
|
||||
valid_types = ["checkpoints", "loras", "embeddings", "vae", "controlnet", "upscalers", "other"]
|
||||
|
||||
if path_type not in valid_types:
|
||||
console.print(f"[red]Error: Invalid type '{path_type}'. Valid: {', '.join(valid_types)}[/red]")
|
||||
if path_type not in VALID_PATH_TYPES:
|
||||
console.print(f"[red]Error: Invalid type '{path_type}'. Valid: {', '.join(VALID_PATH_TYPES)}[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
cfg = load_config()
|
||||
@@ -713,10 +713,7 @@ def config(
|
||||
configured_paths = cfg.get("paths", {})
|
||||
|
||||
for path_str, types in sorted(shown_paths.items(), key=lambda x: x[0]):
|
||||
is_custom = any(
|
||||
path_str == configured_paths.get(k)
|
||||
for k in ["checkpoints", "loras", "embeddings", "vae", "controlnet", "upscalers", "other"]
|
||||
)
|
||||
is_custom = any(path_str == configured_paths.get(k) for k in VALID_PATH_TYPES)
|
||||
marker = " [green](custom)[/green]" if is_custom else " [dim](default)[/dim]"
|
||||
console.print(f" {', '.join(sorted(types))}: {path_str}{marker}")
|
||||
|
||||
|
||||
+51
-1
@@ -26,6 +26,11 @@ GALLERY_DIR = DATA_DIR / "gallery"
|
||||
LEGACY_RC_FILE = Path.home() / ".sftrc"
|
||||
|
||||
# Default download paths by model type (can be overridden in config.toml [paths])
|
||||
#
|
||||
# Note: "DiffusionModel" is not an official CivitAI model_type — CivitAI lumps
|
||||
# UNet-only files (e.g. Flux UNet released separately from CLIP+VAE) under
|
||||
# "Checkpoint". The DiffusionModel entry here exists so users can register a
|
||||
# ComfyUI `diffusion_models/` path and target it manually via `tsr dl -o`.
|
||||
DEFAULT_PATHS: dict[str, Path] = {
|
||||
"Checkpoint": MODELS_DIR / "checkpoints",
|
||||
"LORA": MODELS_DIR / "loras",
|
||||
@@ -34,9 +39,23 @@ DEFAULT_PATHS: dict[str, Path] = {
|
||||
"VAE": MODELS_DIR / "vae",
|
||||
"Controlnet": MODELS_DIR / "controlnet",
|
||||
"Upscaler": MODELS_DIR / "upscalers",
|
||||
"DiffusionModel": MODELS_DIR / "diffusion_models",
|
||||
"Other": MODELS_DIR / "other",
|
||||
}
|
||||
|
||||
# Config-file keys accepted by `tsr config --set-path KEY=PATH`. Single source
|
||||
# of truth shared between the CLI validator and the display-marker logic.
|
||||
VALID_PATH_TYPES: list[str] = [
|
||||
"checkpoints",
|
||||
"loras",
|
||||
"embeddings",
|
||||
"vae",
|
||||
"controlnet",
|
||||
"upscalers",
|
||||
"diffusion_models",
|
||||
"other",
|
||||
]
|
||||
|
||||
CIVITAI_API_BASE = "https://civitai.com/api/v1"
|
||||
CIVITAI_DOWNLOAD_BASE = "https://civitai.com/api/download/models"
|
||||
|
||||
@@ -297,7 +316,8 @@ def get_model_paths() -> dict[str, Path]:
|
||||
config = load_config()
|
||||
paths_config = config.get("paths", {})
|
||||
|
||||
# Map config keys to CivitAI model types
|
||||
# Map config keys to CivitAI model types. "diffusion_models" maps to the
|
||||
# synthetic "DiffusionModel" bucket (see DEFAULT_PATHS for rationale).
|
||||
key_to_types = {
|
||||
"checkpoints": ["Checkpoint"],
|
||||
"loras": ["LORA", "LoCon"],
|
||||
@@ -305,6 +325,7 @@ def get_model_paths() -> dict[str, Path]:
|
||||
"vae": ["VAE"],
|
||||
"controlnet": ["Controlnet"],
|
||||
"upscalers": ["Upscaler"],
|
||||
"diffusion_models": ["DiffusionModel"],
|
||||
"other": ["Other"],
|
||||
}
|
||||
|
||||
@@ -745,6 +766,7 @@ FLUX_UNET_ONLY_PATTERNS: tuple[str, ...] = (
|
||||
"getphatflux", # getphatFLUXReality_v11Softcore.safetensors
|
||||
"moodydesire", # moodyDesireMix_v20PRO.safetensors
|
||||
"fcfluxpony", # fcFluxPonyPerfectBase_fcFluxPerfectBase.safetensors
|
||||
"prototype_", # prototype_v10.safetensors (Flux unet-only, no "flux" in name)
|
||||
)
|
||||
|
||||
|
||||
@@ -753,6 +775,26 @@ def _is_flux_unet_only(name_lower: str) -> bool:
|
||||
return any(p in name_lower for p in FLUX_UNET_ONLY_PATTERNS)
|
||||
|
||||
|
||||
# All-in-one Flux checkpoint filename substrings (case-insensitive). These
|
||||
# checkpoints bundle UNet + CLIP-L + T5 + VAE in a single file (loadable via
|
||||
# CheckpointLoaderSimple), but their filename does NOT contain "flux" so the
|
||||
# generic substring check misses them and they fall through to SDXL defaults
|
||||
# (which fails when the SDXL VAE isn't installed on the target backend).
|
||||
#
|
||||
# Detect via header inspection: keys like `model.diffusion_model.double_blocks.*`
|
||||
# plus bundled `text_encoders.*` and `vae.*` prefixes indicate FLUX all-in-one.
|
||||
# Add new patterns here as we encounter them.
|
||||
FLUX_ALL_IN_ONE_PATTERNS: tuple[str, ...] = (
|
||||
"ultrasense", # ultrasenseInfinity_v10.safetensors
|
||||
"bodyslider", # bodySliderFitness_v10.safetensors
|
||||
)
|
||||
|
||||
|
||||
def _is_flux_all_in_one(name_lower: str) -> bool:
|
||||
"""True if the lowercased filename matches a known all-in-one Flux pattern."""
|
||||
return any(p in name_lower for p in FLUX_ALL_IN_ONE_PATTERNS)
|
||||
|
||||
|
||||
# Flux.2 Klein 9B filename substrings (case-insensitive). These checkpoints are
|
||||
# UNet-only AND require the Flux.2 architecture (Qwen3-8B encoder, Flux2
|
||||
# scheduler). Detection is primarily via base_model field ("Flux.2 Klein"); the
|
||||
@@ -804,6 +846,14 @@ def detect_model_family(model_name: str, base_model: str | None = None) -> str |
|
||||
if _is_flux_unet_only(name_lower):
|
||||
return "flux_unet"
|
||||
|
||||
# All-in-one Flux override for checkpoints whose filename omits "flux"
|
||||
# (e.g. "ultrasenseInfinity_v10.safetensors"). Without this they fall
|
||||
# through to the SDXL default at the bottom of this function and the
|
||||
# generated workflow asks ComfyUI for sdxl_vae.safetensors — which fails
|
||||
# on Flux-only backends like sin.
|
||||
if _is_flux_all_in_one(name_lower):
|
||||
return "flux"
|
||||
|
||||
# Architecture override: filename containing "flux" wins over any base_model
|
||||
# field (handles hybrid models like "FluxPony" that CivitAI tags as "Pony"
|
||||
# but are architecturally Flux and need the Flux workflow).
|
||||
|
||||
Reference in New Issue
Block a user