From cd8cea67ef88c42006f837c8f780b177e1a82326 Mon Sep 17 00:00:00 2001 From: aladac Date: Mon, 18 May 2026 19:29:46 +0200 Subject: [PATCH] feat(config): register diffusion_models as a first-class path key MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 -o /path/to/diffusion_models` for UNet-only Flux downloads, no more symlink workarounds. --- tensors/cli.py | 11 ++++------ tensors/config.py | 52 ++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 55 insertions(+), 8 deletions(-) diff --git a/tensors/cli.py b/tensors/cli.py index 4c3e957..c5fcee1 100644 --- a/tensors/cli.py +++ b/tensors/cli.py @@ -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}") diff --git a/tensors/config.py b/tensors/config.py index ffd17ef..19f46f5 100644 --- a/tensors/config.py +++ b/tensors/config.py @@ -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).