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,
|
COMFYUI_DEFAULT_WIDTH,
|
||||||
CONFIG_FILE,
|
CONFIG_FILE,
|
||||||
MODEL_FAMILY_DEFAULTS,
|
MODEL_FAMILY_DEFAULTS,
|
||||||
|
VALID_PATH_TYPES,
|
||||||
BaseModel,
|
BaseModel,
|
||||||
CommercialUse,
|
CommercialUse,
|
||||||
ModelType,
|
ModelType,
|
||||||
@@ -673,10 +674,9 @@ def config(
|
|||||||
|
|
||||||
path_type, path_value = set_path.split("=", 1)
|
path_type, path_value = set_path.split("=", 1)
|
||||||
path_type = path_type.lower().strip()
|
path_type = path_type.lower().strip()
|
||||||
valid_types = ["checkpoints", "loras", "embeddings", "vae", "controlnet", "upscalers", "other"]
|
|
||||||
|
|
||||||
if path_type not in valid_types:
|
if path_type not in VALID_PATH_TYPES:
|
||||||
console.print(f"[red]Error: Invalid type '{path_type}'. Valid: {', '.join(valid_types)}[/red]")
|
console.print(f"[red]Error: Invalid type '{path_type}'. Valid: {', '.join(VALID_PATH_TYPES)}[/red]")
|
||||||
raise typer.Exit(1)
|
raise typer.Exit(1)
|
||||||
|
|
||||||
cfg = load_config()
|
cfg = load_config()
|
||||||
@@ -713,10 +713,7 @@ def config(
|
|||||||
configured_paths = cfg.get("paths", {})
|
configured_paths = cfg.get("paths", {})
|
||||||
|
|
||||||
for path_str, types in sorted(shown_paths.items(), key=lambda x: x[0]):
|
for path_str, types in sorted(shown_paths.items(), key=lambda x: x[0]):
|
||||||
is_custom = any(
|
is_custom = any(path_str == configured_paths.get(k) for k in VALID_PATH_TYPES)
|
||||||
path_str == configured_paths.get(k)
|
|
||||||
for k in ["checkpoints", "loras", "embeddings", "vae", "controlnet", "upscalers", "other"]
|
|
||||||
)
|
|
||||||
marker = " [green](custom)[/green]" if is_custom else " [dim](default)[/dim]"
|
marker = " [green](custom)[/green]" if is_custom else " [dim](default)[/dim]"
|
||||||
console.print(f" {', '.join(sorted(types))}: {path_str}{marker}")
|
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"
|
LEGACY_RC_FILE = Path.home() / ".sftrc"
|
||||||
|
|
||||||
# Default download paths by model type (can be overridden in config.toml [paths])
|
# 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] = {
|
DEFAULT_PATHS: dict[str, Path] = {
|
||||||
"Checkpoint": MODELS_DIR / "checkpoints",
|
"Checkpoint": MODELS_DIR / "checkpoints",
|
||||||
"LORA": MODELS_DIR / "loras",
|
"LORA": MODELS_DIR / "loras",
|
||||||
@@ -34,9 +39,23 @@ DEFAULT_PATHS: dict[str, Path] = {
|
|||||||
"VAE": MODELS_DIR / "vae",
|
"VAE": MODELS_DIR / "vae",
|
||||||
"Controlnet": MODELS_DIR / "controlnet",
|
"Controlnet": MODELS_DIR / "controlnet",
|
||||||
"Upscaler": MODELS_DIR / "upscalers",
|
"Upscaler": MODELS_DIR / "upscalers",
|
||||||
|
"DiffusionModel": MODELS_DIR / "diffusion_models",
|
||||||
"Other": MODELS_DIR / "other",
|
"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_API_BASE = "https://civitai.com/api/v1"
|
||||||
CIVITAI_DOWNLOAD_BASE = "https://civitai.com/api/download/models"
|
CIVITAI_DOWNLOAD_BASE = "https://civitai.com/api/download/models"
|
||||||
|
|
||||||
@@ -297,7 +316,8 @@ def get_model_paths() -> dict[str, Path]:
|
|||||||
config = load_config()
|
config = load_config()
|
||||||
paths_config = config.get("paths", {})
|
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 = {
|
key_to_types = {
|
||||||
"checkpoints": ["Checkpoint"],
|
"checkpoints": ["Checkpoint"],
|
||||||
"loras": ["LORA", "LoCon"],
|
"loras": ["LORA", "LoCon"],
|
||||||
@@ -305,6 +325,7 @@ def get_model_paths() -> dict[str, Path]:
|
|||||||
"vae": ["VAE"],
|
"vae": ["VAE"],
|
||||||
"controlnet": ["Controlnet"],
|
"controlnet": ["Controlnet"],
|
||||||
"upscalers": ["Upscaler"],
|
"upscalers": ["Upscaler"],
|
||||||
|
"diffusion_models": ["DiffusionModel"],
|
||||||
"other": ["Other"],
|
"other": ["Other"],
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -745,6 +766,7 @@ FLUX_UNET_ONLY_PATTERNS: tuple[str, ...] = (
|
|||||||
"getphatflux", # getphatFLUXReality_v11Softcore.safetensors
|
"getphatflux", # getphatFLUXReality_v11Softcore.safetensors
|
||||||
"moodydesire", # moodyDesireMix_v20PRO.safetensors
|
"moodydesire", # moodyDesireMix_v20PRO.safetensors
|
||||||
"fcfluxpony", # fcFluxPonyPerfectBase_fcFluxPerfectBase.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)
|
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
|
# Flux.2 Klein 9B filename substrings (case-insensitive). These checkpoints are
|
||||||
# UNet-only AND require the Flux.2 architecture (Qwen3-8B encoder, Flux2
|
# UNet-only AND require the Flux.2 architecture (Qwen3-8B encoder, Flux2
|
||||||
# scheduler). Detection is primarily via base_model field ("Flux.2 Klein"); the
|
# 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):
|
if _is_flux_unet_only(name_lower):
|
||||||
return "flux_unet"
|
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
|
# Architecture override: filename containing "flux" wins over any base_model
|
||||||
# field (handles hybrid models like "FluxPony" that CivitAI tags as "Pony"
|
# field (handles hybrid models like "FluxPony" that CivitAI tags as "Pony"
|
||||||
# but are architecturally Flux and need the Flux workflow).
|
# but are architecturally Flux and need the Flux workflow).
|
||||||
|
|||||||
Reference in New Issue
Block a user