mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-21 11:11:26 -03:00
Compare commits
5 Commits
d43ab6e32f
...
c85b6b64a1
| Author | SHA1 | Date | |
|---|---|---|---|
| c85b6b64a1 | |||
| 34c87d4934 | |||
| 93472e5d67 | |||
| ae185ee714 | |||
| 795036275a |
+10
@@ -3,6 +3,8 @@ try: # pragma: no cover - import fallback for pytest collection
|
||||
from .py.nodes.lora_loader import LoraLoaderLM, LoraTextLoaderLM
|
||||
from .py.nodes.checkpoint_loader import CheckpointLoaderLM
|
||||
from .py.nodes.unet_loader import UNETLoaderLM
|
||||
from .py.nodes.random_checkpoint_loader import RandomCheckpointLoaderLM
|
||||
from .py.nodes.random_unet_loader import RandomUNETLoaderLM
|
||||
from .py.nodes.trigger_word_toggle import TriggerWordToggleLM
|
||||
from .py.nodes.prompt import PromptLM
|
||||
from .py.nodes.text import TextLM
|
||||
@@ -40,6 +42,12 @@ except (
|
||||
"py.nodes.checkpoint_loader"
|
||||
).CheckpointLoaderLM
|
||||
UNETLoaderLM = importlib.import_module("py.nodes.unet_loader").UNETLoaderLM
|
||||
RandomCheckpointLoaderLM = importlib.import_module(
|
||||
"py.nodes.random_checkpoint_loader"
|
||||
).RandomCheckpointLoaderLM
|
||||
RandomUNETLoaderLM = importlib.import_module(
|
||||
"py.nodes.random_unet_loader"
|
||||
).RandomUNETLoaderLM
|
||||
TriggerWordToggleLM = importlib.import_module(
|
||||
"py.nodes.trigger_word_toggle"
|
||||
).TriggerWordToggleLM
|
||||
@@ -79,6 +87,8 @@ NODE_CLASS_MAPPINGS = {
|
||||
LoraTextLoaderLM.NAME: LoraTextLoaderLM,
|
||||
CheckpointLoaderLM.NAME: CheckpointLoaderLM,
|
||||
UNETLoaderLM.NAME: UNETLoaderLM,
|
||||
RandomCheckpointLoaderLM.NAME: RandomCheckpointLoaderLM,
|
||||
RandomUNETLoaderLM.NAME: RandomUNETLoaderLM,
|
||||
TriggerWordToggleLM.NAME: TriggerWordToggleLM,
|
||||
LoraStackerLM.NAME: LoraStackerLM,
|
||||
LoraStackCombinerLM.NAME: LoraStackCombinerLM,
|
||||
|
||||
+3
-1
@@ -920,7 +920,9 @@
|
||||
"dateAsc": "Älteste",
|
||||
"lorasCount": "LoRA-Anzahl",
|
||||
"lorasCountDesc": "Meiste",
|
||||
"lorasCountAsc": "Wenigste"
|
||||
"lorasCountAsc": "Wenigste",
|
||||
"opened": "Zuletzt geöffnet",
|
||||
"openedDesc": "Zuletzt geöffnet"
|
||||
},
|
||||
"refresh": {
|
||||
"title": "Rezeptliste aktualisieren",
|
||||
|
||||
+3
-1
@@ -920,7 +920,9 @@
|
||||
"dateAsc": "Oldest",
|
||||
"lorasCount": "LoRA Count",
|
||||
"lorasCountDesc": "Most",
|
||||
"lorasCountAsc": "Least"
|
||||
"lorasCountAsc": "Least",
|
||||
"opened": "Recently Opened",
|
||||
"openedDesc": "Recently opened"
|
||||
},
|
||||
"refresh": {
|
||||
"title": "Refresh recipe list",
|
||||
|
||||
+3
-1
@@ -920,7 +920,9 @@
|
||||
"dateAsc": "Más antiguo",
|
||||
"lorasCount": "Cant. de LoRAs",
|
||||
"lorasCountDesc": "Más",
|
||||
"lorasCountAsc": "Menos"
|
||||
"lorasCountAsc": "Menos",
|
||||
"opened": "Abiertos recientemente",
|
||||
"openedDesc": "Abiertos recientemente"
|
||||
},
|
||||
"refresh": {
|
||||
"title": "Actualizar lista de recetas",
|
||||
|
||||
+3
-1
@@ -920,7 +920,9 @@
|
||||
"dateAsc": "Plus ancien",
|
||||
"lorasCount": "Nombre de LoRAs",
|
||||
"lorasCountDesc": "Plus",
|
||||
"lorasCountAsc": "Moins"
|
||||
"lorasCountAsc": "Moins",
|
||||
"opened": "Récemment ouverts",
|
||||
"openedDesc": "Récemment ouverts"
|
||||
},
|
||||
"refresh": {
|
||||
"title": "Actualiser la liste des recipes",
|
||||
|
||||
+3
-1
@@ -920,7 +920,9 @@
|
||||
"dateAsc": "הכי ישן",
|
||||
"lorasCount": "מספר LoRAs",
|
||||
"lorasCountDesc": "הכי הרבה",
|
||||
"lorasCountAsc": "הכי פחות"
|
||||
"lorasCountAsc": "הכי פחות",
|
||||
"opened": "נפתחו לאחרונה",
|
||||
"openedDesc": "נפתחו לאחרונה"
|
||||
},
|
||||
"refresh": {
|
||||
"title": "רענן רשימת מתכונים",
|
||||
|
||||
+3
-1
@@ -920,7 +920,9 @@
|
||||
"dateAsc": "古い順",
|
||||
"lorasCount": "LoRA数",
|
||||
"lorasCountDesc": "多い順",
|
||||
"lorasCountAsc": "少ない順"
|
||||
"lorasCountAsc": "少ない順",
|
||||
"opened": "最近開いた",
|
||||
"openedDesc": "最近開いた"
|
||||
},
|
||||
"refresh": {
|
||||
"title": "レシピリストを更新",
|
||||
|
||||
+3
-1
@@ -920,7 +920,9 @@
|
||||
"dateAsc": "오래된순",
|
||||
"lorasCount": "LoRA 수",
|
||||
"lorasCountDesc": "많은순",
|
||||
"lorasCountAsc": "적은순"
|
||||
"lorasCountAsc": "적은순",
|
||||
"opened": "최근에 연",
|
||||
"openedDesc": "최근에 연"
|
||||
},
|
||||
"refresh": {
|
||||
"title": "레시피 목록 새로고침",
|
||||
|
||||
+3
-1
@@ -920,7 +920,9 @@
|
||||
"dateAsc": "Сначала старые",
|
||||
"lorasCount": "Кол-во LoRA",
|
||||
"lorasCountDesc": "Больше всего",
|
||||
"lorasCountAsc": "Меньше всего"
|
||||
"lorasCountAsc": "Меньше всего",
|
||||
"opened": "Недавно открытые",
|
||||
"openedDesc": "Недавно открытые"
|
||||
},
|
||||
"refresh": {
|
||||
"title": "Обновить список рецептов",
|
||||
|
||||
+3
-1
@@ -920,7 +920,9 @@
|
||||
"dateAsc": "最早",
|
||||
"lorasCount": "LoRA 数量",
|
||||
"lorasCountDesc": "最多",
|
||||
"lorasCountAsc": "最少"
|
||||
"lorasCountAsc": "最少",
|
||||
"opened": "最近打开",
|
||||
"openedDesc": "最近打开"
|
||||
},
|
||||
"refresh": {
|
||||
"title": "刷新配方列表",
|
||||
|
||||
+3
-1
@@ -920,7 +920,9 @@
|
||||
"dateAsc": "最舊",
|
||||
"lorasCount": "LoRA 數量",
|
||||
"lorasCountDesc": "最多",
|
||||
"lorasCountAsc": "最少"
|
||||
"lorasCountAsc": "最少",
|
||||
"opened": "最近開啟",
|
||||
"openedDesc": "最近開啟"
|
||||
},
|
||||
"refresh": {
|
||||
"title": "重新整理配方列表",
|
||||
|
||||
@@ -0,0 +1,214 @@
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
from typing import Any, List, Optional, Tuple
|
||||
import comfy.sd # pyright: ignore[reportMissingImports]
|
||||
import folder_paths # pyright: ignore[reportMissingImports]
|
||||
from ..utils.utils import get_checkpoint_info_absolute, _format_model_name_for_comfyui
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RandomCheckpointLoaderLM:
|
||||
"""Checkpoint Loader that can randomly pick a checkpoint from the pool
|
||||
|
||||
Loads checkpoints from both standard ComfyUI folders and LoRA Manager's
|
||||
extra folder paths. When select_at_random is enabled, ignores ckpt_name
|
||||
and picks a random checkpoint (optionally filtered by base_model) on
|
||||
every run.
|
||||
"""
|
||||
|
||||
NAME = "Random Checkpoint Loader (LoraManager)"
|
||||
CATEGORY = "Lora Manager/loaders"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
# Get list of checkpoint names from scanner (includes extra folder paths)
|
||||
checkpoint_names = cls._get_checkpoint_names()
|
||||
base_models = cls._get_available_base_models()
|
||||
return {
|
||||
"required": {
|
||||
"ckpt_name": (
|
||||
checkpoint_names,
|
||||
{"tooltip": "The name of the checkpoint (model) to load."},
|
||||
),
|
||||
"select_at_random": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": (
|
||||
"Ignore ckpt_name and pick a random checkpoint from the "
|
||||
"pool (optionally filtered by base_model) on every run."
|
||||
),
|
||||
},
|
||||
),
|
||||
"base_model": (
|
||||
base_models,
|
||||
{
|
||||
"default": "Any",
|
||||
"tooltip": "Restrict random selection to this base model. 'Any' uses the full pool.",
|
||||
},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE", "STRING")
|
||||
RETURN_NAMES = ("MODEL", "CLIP", "VAE", "model_name")
|
||||
OUTPUT_TOOLTIPS = (
|
||||
"The model used for denoising latents.",
|
||||
"The CLIP model used for encoding text prompts.",
|
||||
"The VAE model used for encoding and decoding images to and from latent space.",
|
||||
"The name of the checkpoint that was loaded (useful when select_at_random is enabled).",
|
||||
)
|
||||
FUNCTION = "load_checkpoint"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, ckpt_name, select_at_random=False, base_model="Any"):
|
||||
# Force re-execution on every run while randomizing, since the widget
|
||||
# values themselves don't change between queue runs.
|
||||
if select_at_random:
|
||||
return float("nan")
|
||||
return ckpt_name
|
||||
|
||||
@staticmethod
|
||||
def _run_async(coro_fn):
|
||||
"""Run an async fetcher, handling the case where an event loop is already running."""
|
||||
import asyncio
|
||||
|
||||
try:
|
||||
asyncio.get_running_loop()
|
||||
import concurrent.futures
|
||||
|
||||
def run_in_thread():
|
||||
new_loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(new_loop)
|
||||
try:
|
||||
return new_loop.run_until_complete(coro_fn())
|
||||
finally:
|
||||
new_loop.close()
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future = executor.submit(run_in_thread)
|
||||
return future.result()
|
||||
except RuntimeError:
|
||||
return asyncio.run(coro_fn())
|
||||
|
||||
@classmethod
|
||||
def _get_checkpoint_names(cls, base_model: Optional[str] = None) -> List[str]:
|
||||
"""Get list of checkpoint names from scanner cache in ComfyUI format (relative path with extension)
|
||||
|
||||
Args:
|
||||
base_model: If given (and not "Any"), only include checkpoints matching this base model.
|
||||
"""
|
||||
try:
|
||||
from ..services.service_registry import ServiceRegistry
|
||||
|
||||
async def _get_names():
|
||||
scanner = await ServiceRegistry.get_checkpoint_scanner()
|
||||
cache = await scanner.get_cached_data()
|
||||
|
||||
# Get all model roots for calculating relative paths
|
||||
model_roots = scanner.get_model_roots()
|
||||
|
||||
# Filter only checkpoint type (not diffusion_model) and format names
|
||||
names = []
|
||||
for item in cache.raw_data:
|
||||
if item.get("sub_type") != "checkpoint":
|
||||
continue
|
||||
if (
|
||||
base_model
|
||||
and base_model != "Any"
|
||||
and item.get("base_model") != base_model
|
||||
):
|
||||
continue
|
||||
file_path = item.get("file_path", "")
|
||||
# Only offer models that still exist on disk so ComfyUI
|
||||
# flags missing checkpoints at queue time via
|
||||
# "value not in list" (the scanner cache can be stale).
|
||||
if file_path and os.path.exists(file_path):
|
||||
# Format using relative path with OS-native separator
|
||||
formatted_name = _format_model_name_for_comfyui(
|
||||
file_path, model_roots
|
||||
)
|
||||
if formatted_name:
|
||||
names.append(formatted_name)
|
||||
|
||||
return sorted(names)
|
||||
|
||||
return cls._run_async(_get_names)
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting checkpoint names: {e}")
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def _get_available_base_models(cls) -> List[str]:
|
||||
"""Get distinct base_model values present among indexed checkpoints, for the random-selection filter."""
|
||||
try:
|
||||
from ..services.service_registry import ServiceRegistry
|
||||
|
||||
async def _get_base_models():
|
||||
scanner = await ServiceRegistry.get_checkpoint_scanner()
|
||||
cache = await scanner.get_cached_data()
|
||||
|
||||
base_models = set()
|
||||
for item in cache.raw_data:
|
||||
if item.get("sub_type") != "checkpoint":
|
||||
continue
|
||||
base_model = item.get("base_model")
|
||||
file_path = item.get("file_path", "")
|
||||
if base_model and file_path and os.path.exists(file_path):
|
||||
base_models.add(base_model)
|
||||
|
||||
return sorted(base_models)
|
||||
|
||||
return ["Any"] + cls._run_async(_get_base_models)
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting available base models: {e}")
|
||||
return ["Any"]
|
||||
|
||||
def load_checkpoint(
|
||||
self,
|
||||
ckpt_name: str,
|
||||
select_at_random: bool = False,
|
||||
base_model: str = "Any",
|
||||
) -> Tuple[Any, Any, Any, str]:
|
||||
"""Load a checkpoint by name, supporting extra folder paths
|
||||
|
||||
Args:
|
||||
ckpt_name: The name of the checkpoint to load (relative path with extension)
|
||||
select_at_random: If True, ignore ckpt_name and pick randomly from the pool
|
||||
base_model: Restricts random selection to this base model ("Any" = no filter)
|
||||
|
||||
Returns:
|
||||
Tuple of (MODEL, CLIP, VAE, model_name)
|
||||
"""
|
||||
if select_at_random:
|
||||
pool = self._get_checkpoint_names(base_model)
|
||||
if not pool:
|
||||
raise FileNotFoundError(
|
||||
f"No checkpoints found for base model '{base_model}'. "
|
||||
"Pick a different base model or disable 'select_at_random'."
|
||||
)
|
||||
ckpt_name = random.choice(pool)
|
||||
logger.info(
|
||||
f"[RandomCheckpointLoaderLM] Randomly selected checkpoint: {ckpt_name}"
|
||||
)
|
||||
|
||||
# Get absolute path from cache using ComfyUI-style name
|
||||
ckpt_path, metadata = get_checkpoint_info_absolute(ckpt_name)
|
||||
|
||||
if metadata is None:
|
||||
raise FileNotFoundError(
|
||||
f"Checkpoint '{ckpt_name}' not found in LoRA Manager cache. "
|
||||
"Make sure the checkpoint is indexed and try again."
|
||||
)
|
||||
|
||||
# Load regular checkpoint using ComfyUI's API
|
||||
logger.info(f"Loading checkpoint from: {ckpt_path}")
|
||||
out = comfy.sd.load_checkpoint_guess_config(
|
||||
ckpt_path,
|
||||
output_vae=True,
|
||||
output_clip=True,
|
||||
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
||||
)
|
||||
return out[:3] + (ckpt_name,)
|
||||
@@ -0,0 +1,326 @@
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
from typing import Any, List, Optional, Tuple
|
||||
import comfy.sd # pyright: ignore[reportMissingImports]
|
||||
from ..utils.utils import get_checkpoint_info_absolute, _format_model_name_for_comfyui
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _reload_gguf_unet(
|
||||
unet_path: str, weight_dtype: str, disable_dynamic: bool = False
|
||||
) -> object:
|
||||
"""Reload a GGUF diffusion model from disk (cached_patcher_init factory).
|
||||
|
||||
Mirrors the GGUF branch of RandomUNETLoaderLM.load_unet so ModelPatcher
|
||||
deepclone/dynamic machinery can rebuild GGUF models with the correct
|
||||
GGMLOps. ``disable_dynamic`` is accepted for signature compatibility
|
||||
with core ComfyUI loaders.
|
||||
"""
|
||||
loader = RandomUNETLoaderLM()
|
||||
model, _unet_name = loader._load_gguf_unet(unet_path, unet_path, weight_dtype)
|
||||
return model
|
||||
|
||||
|
||||
class RandomUNETLoaderLM:
|
||||
"""UNET Loader that can randomly pick a diffusion model from the pool
|
||||
|
||||
Loads diffusion models/UNets from both standard ComfyUI folders and LoRA
|
||||
Manager's extra folder paths. Supports both regular diffusion models and
|
||||
GGUF format models. When select_at_random is enabled, ignores unet_name
|
||||
and picks a random diffusion model (optionally filtered by base_model)
|
||||
on every run.
|
||||
"""
|
||||
|
||||
NAME = "Random Unet Loader (LoraManager)"
|
||||
CATEGORY = "Lora Manager/loaders"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
# Get list of unet names from scanner (includes extra folder paths)
|
||||
unet_names = cls._get_unet_names()
|
||||
base_models = cls._get_available_base_models()
|
||||
return {
|
||||
"required": {
|
||||
"unet_name": (
|
||||
unet_names,
|
||||
{"tooltip": "The name of the diffusion model to load."},
|
||||
),
|
||||
"weight_dtype": (
|
||||
["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],
|
||||
{"tooltip": "The dtype to use for the model weights."},
|
||||
),
|
||||
"select_at_random": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": (
|
||||
"Ignore unet_name and pick a random diffusion model from "
|
||||
"the pool (optionally filtered by base_model) on every run."
|
||||
),
|
||||
},
|
||||
),
|
||||
"base_model": (
|
||||
base_models,
|
||||
{
|
||||
"default": "Any",
|
||||
"tooltip": "Restrict random selection to this base model. 'Any' uses the full pool.",
|
||||
},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "STRING")
|
||||
RETURN_NAMES = ("MODEL", "model_name")
|
||||
OUTPUT_TOOLTIPS = (
|
||||
"The model used for denoising latents.",
|
||||
"The name of the diffusion model that was loaded (useful when select_at_random is enabled).",
|
||||
)
|
||||
FUNCTION = "load_unet"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(
|
||||
cls, unet_name, weight_dtype, select_at_random=False, base_model="Any"
|
||||
):
|
||||
# Force re-execution on every run while randomizing, since the widget
|
||||
# values themselves don't change between queue runs.
|
||||
if select_at_random:
|
||||
return float("nan")
|
||||
return unet_name
|
||||
|
||||
@staticmethod
|
||||
def _run_async(coro_fn):
|
||||
"""Run an async fetcher, handling the case where an event loop is already running."""
|
||||
import asyncio
|
||||
|
||||
try:
|
||||
asyncio.get_running_loop()
|
||||
import concurrent.futures
|
||||
|
||||
def run_in_thread():
|
||||
new_loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(new_loop)
|
||||
try:
|
||||
return new_loop.run_until_complete(coro_fn())
|
||||
finally:
|
||||
new_loop.close()
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future = executor.submit(run_in_thread)
|
||||
return future.result()
|
||||
except RuntimeError:
|
||||
return asyncio.run(coro_fn())
|
||||
|
||||
@classmethod
|
||||
def _get_unet_names(cls, base_model: Optional[str] = None) -> List[str]:
|
||||
"""Get list of diffusion model names from scanner cache in ComfyUI format (relative path with extension)
|
||||
|
||||
Args:
|
||||
base_model: If given (and not "Any"), only include models matching this base model.
|
||||
"""
|
||||
try:
|
||||
from ..services.service_registry import ServiceRegistry
|
||||
|
||||
async def _get_names():
|
||||
scanner = await ServiceRegistry.get_checkpoint_scanner()
|
||||
cache = await scanner.get_cached_data()
|
||||
|
||||
# Get all model roots for calculating relative paths
|
||||
model_roots = scanner.get_model_roots()
|
||||
|
||||
# Filter only diffusion_model type and format names
|
||||
names = []
|
||||
for item in cache.raw_data:
|
||||
if item.get("sub_type") != "diffusion_model":
|
||||
continue
|
||||
if (
|
||||
base_model
|
||||
and base_model != "Any"
|
||||
and item.get("base_model") != base_model
|
||||
):
|
||||
continue
|
||||
file_path = item.get("file_path", "")
|
||||
# Only offer models that still exist on disk so ComfyUI
|
||||
# flags missing diffusion models at queue time via
|
||||
# "value not in list" (the scanner cache can be stale).
|
||||
if file_path and os.path.exists(file_path):
|
||||
# Format using relative path with OS-native separator
|
||||
formatted_name = _format_model_name_for_comfyui(
|
||||
file_path, model_roots
|
||||
)
|
||||
if formatted_name:
|
||||
names.append(formatted_name)
|
||||
|
||||
return sorted(names)
|
||||
|
||||
return cls._run_async(_get_names)
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting unet names: {e}")
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def _get_available_base_models(cls) -> List[str]:
|
||||
"""Get distinct base_model values present among indexed diffusion models, for the random-selection filter."""
|
||||
try:
|
||||
from ..services.service_registry import ServiceRegistry
|
||||
|
||||
async def _get_base_models():
|
||||
scanner = await ServiceRegistry.get_checkpoint_scanner()
|
||||
cache = await scanner.get_cached_data()
|
||||
|
||||
base_models = set()
|
||||
for item in cache.raw_data:
|
||||
if item.get("sub_type") != "diffusion_model":
|
||||
continue
|
||||
base_model = item.get("base_model")
|
||||
file_path = item.get("file_path", "")
|
||||
if base_model and file_path and os.path.exists(file_path):
|
||||
base_models.add(base_model)
|
||||
|
||||
return sorted(base_models)
|
||||
|
||||
return ["Any"] + cls._run_async(_get_base_models)
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting available base models: {e}")
|
||||
return ["Any"]
|
||||
|
||||
def load_unet(
|
||||
self,
|
||||
unet_name: str,
|
||||
weight_dtype: str,
|
||||
select_at_random: bool = False,
|
||||
base_model: str = "Any",
|
||||
) -> Tuple[Any, ...]:
|
||||
"""Load a diffusion model by name, supporting extra folder paths
|
||||
|
||||
Args:
|
||||
unet_name: The name of the diffusion model to load (relative path with extension)
|
||||
weight_dtype: The dtype to use for model weights
|
||||
select_at_random: If True, ignore unet_name and pick randomly from the pool
|
||||
base_model: Restricts random selection to this base model ("Any" = no filter)
|
||||
|
||||
Returns:
|
||||
Tuple of (MODEL, model_name)
|
||||
"""
|
||||
import torch
|
||||
|
||||
if select_at_random:
|
||||
pool = self._get_unet_names(base_model)
|
||||
if not pool:
|
||||
raise FileNotFoundError(
|
||||
f"No diffusion models found for base model '{base_model}'. "
|
||||
"Pick a different base model or disable 'select_at_random'."
|
||||
)
|
||||
unet_name = random.choice(pool)
|
||||
logger.info(
|
||||
f"[RandomUNETLoaderLM] Randomly selected diffusion model: {unet_name}"
|
||||
)
|
||||
|
||||
# Get absolute path from cache using ComfyUI-style name
|
||||
unet_path, metadata = get_checkpoint_info_absolute(unet_name)
|
||||
|
||||
if metadata is None:
|
||||
raise FileNotFoundError(
|
||||
f"Diffusion model '{unet_name}' not found in LoRA Manager cache. "
|
||||
"Make sure the model is indexed and try again."
|
||||
)
|
||||
|
||||
# Check if it's a GGUF model
|
||||
if unet_path.endswith(".gguf"):
|
||||
return self._load_gguf_unet(unet_path, unet_name, weight_dtype)
|
||||
|
||||
# Load regular diffusion model using ComfyUI's API
|
||||
logger.info(f"Loading diffusion model from: {unet_path}")
|
||||
|
||||
# Build model options based on weight_dtype
|
||||
model_options = {}
|
||||
if weight_dtype == "fp8_e4m3fn":
|
||||
model_options["dtype"] = torch.float8_e4m3fn
|
||||
elif weight_dtype == "fp8_e4m3fn_fast":
|
||||
model_options["dtype"] = torch.float8_e4m3fn
|
||||
model_options["fp8_optimizations"] = True
|
||||
elif weight_dtype == "fp8_e5m2":
|
||||
model_options["dtype"] = torch.float8_e5m2
|
||||
|
||||
model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options)
|
||||
return (model, unet_name)
|
||||
|
||||
def _load_gguf_unet(
|
||||
self, unet_path: str, unet_name: str, weight_dtype: str
|
||||
) -> Tuple[Any, ...]:
|
||||
"""Load a GGUF format diffusion model
|
||||
|
||||
Args:
|
||||
unet_path: Absolute path to the GGUF file
|
||||
unet_name: Name of the model for error messages
|
||||
weight_dtype: The dtype to use for model weights
|
||||
|
||||
Returns:
|
||||
Tuple of (MODEL, model_name)
|
||||
"""
|
||||
import torch
|
||||
from .gguf_import_helper import get_gguf_modules
|
||||
|
||||
# Get ComfyUI-GGUF modules using helper (handles various import scenarios)
|
||||
try:
|
||||
loader_module, ops_module, nodes_module = get_gguf_modules()
|
||||
gguf_sd_loader = getattr(loader_module, "gguf_sd_loader")
|
||||
GGMLOps = getattr(ops_module, "GGMLOps")
|
||||
GGUFModelPatcher = getattr(nodes_module, "GGUFModelPatcher")
|
||||
except RuntimeError as e:
|
||||
raise RuntimeError(f"Cannot load GGUF model '{unet_name}'. {str(e)}")
|
||||
|
||||
logger.info(f"Loading GGUF diffusion model from: {unet_path}")
|
||||
|
||||
try:
|
||||
# Load GGUF state dict
|
||||
sd, extra = gguf_sd_loader(unet_path)
|
||||
|
||||
# Prepare kwargs for metadata if supported
|
||||
kwargs = {}
|
||||
import inspect
|
||||
|
||||
valid_params = inspect.signature(
|
||||
comfy.sd.load_diffusion_model_state_dict
|
||||
).parameters
|
||||
if "metadata" in valid_params:
|
||||
kwargs["metadata"] = extra.get("metadata", {})
|
||||
|
||||
# Setup custom operations with GGUF support
|
||||
ops = GGMLOps()
|
||||
|
||||
# Handle weight_dtype for GGUF models
|
||||
if weight_dtype in ("default", None):
|
||||
ops.Linear.dequant_dtype = None
|
||||
elif weight_dtype in ["target"]:
|
||||
ops.Linear.dequant_dtype = weight_dtype
|
||||
else:
|
||||
ops.Linear.dequant_dtype = getattr(torch, weight_dtype, None)
|
||||
|
||||
# Load the model
|
||||
model = comfy.sd.load_diffusion_model_state_dict(
|
||||
sd, model_options={"custom_operations": ops}, **kwargs
|
||||
)
|
||||
|
||||
if model is None:
|
||||
raise RuntimeError(
|
||||
f"Could not detect model type for GGUF diffusion model: {unet_path}"
|
||||
)
|
||||
|
||||
# Wrap with GGUFModelPatcher
|
||||
model = GGUFModelPatcher.clone(model)
|
||||
|
||||
# Register a reload factory so the MODEL carries its source path
|
||||
# (cached_patcher_init) like core ComfyUI loaders do — required
|
||||
# for model-name extraction downstream and for ModelPatcher
|
||||
# deepclone/dynamic machinery.
|
||||
model.cached_patcher_init = (_reload_gguf_unet, (unet_path, weight_dtype))
|
||||
|
||||
return (model, unet_name)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading GGUF diffusion model '{unet_name}': {e}")
|
||||
raise RuntimeError(
|
||||
f"Failed to load GGUF diffusion model '{unet_name}': {str(e)}"
|
||||
)
|
||||
@@ -34,6 +34,7 @@ from ...utils.civitai_utils import (
|
||||
)
|
||||
from ...utils.constants import NSFW_LEVELS
|
||||
from ...utils.exif_utils import ExifUtils
|
||||
from ...utils.recipe_open_stats import RecipeOpenStats
|
||||
from ...recipes.merger import GenParamsMerger
|
||||
from ...recipes.enrichment import RecipeEnricher
|
||||
from ...services.websocket_manager import ws_manager as default_ws_manager
|
||||
@@ -98,6 +99,7 @@ class RecipeHandlerSet:
|
||||
"download_shared_recipe": self.sharing.download_shared_recipe,
|
||||
"get_recipe_syntax": self.query.get_recipe_syntax,
|
||||
"update_recipe": self.management.update_recipe,
|
||||
"record_recipe_open": self.management.record_recipe_open,
|
||||
"reconnect_lora": self.management.reconnect_lora,
|
||||
"find_duplicates": self.query.find_duplicates,
|
||||
"move_recipes_bulk": self.management.move_recipes_bulk,
|
||||
@@ -1458,6 +1460,33 @@ class RecipeManagementHandler:
|
||||
self._logger.error("Error updating recipe: %s", exc, exc_info=True)
|
||||
return web.json_response({"error": str(exc)}, status=500)
|
||||
|
||||
async def record_recipe_open(self, request: web.Request) -> web.Response:
|
||||
"""Record that a recipe's detail modal was opened.
|
||||
|
||||
Lightweight fire-and-forget endpoint backing the "Recently Opened"
|
||||
sort. It only writes the timestamp into the separate open-stats file
|
||||
— recipe JSON and EXIF are never touched.
|
||||
"""
|
||||
try:
|
||||
await self._ensure_dependencies_ready()
|
||||
recipe_scanner = self._recipe_scanner_getter()
|
||||
if recipe_scanner is None:
|
||||
raise RuntimeError("Recipe scanner unavailable")
|
||||
|
||||
recipe_id = request.match_info["recipe_id"]
|
||||
# Skip recording opens for recipes the scanner no longer knows.
|
||||
recipe_json_path = await recipe_scanner.get_recipe_json_path(recipe_id)
|
||||
if not recipe_json_path:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Recipe not found"}, status=404
|
||||
)
|
||||
|
||||
RecipeOpenStats().record_open(recipe_id)
|
||||
return web.json_response({"success": True})
|
||||
except Exception as exc:
|
||||
self._logger.error("Error recording recipe open: %s", exc, exc_info=True)
|
||||
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
||||
|
||||
async def move_recipe(self, request: web.Request) -> web.Response:
|
||||
try:
|
||||
await self._ensure_dependencies_ready()
|
||||
|
||||
@@ -43,6 +43,9 @@ ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
|
||||
),
|
||||
RouteDefinition("GET", "/api/lm/recipe/{recipe_id}/syntax", "get_recipe_syntax"),
|
||||
RouteDefinition("PUT", "/api/lm/recipe/{recipe_id}/update", "update_recipe"),
|
||||
RouteDefinition(
|
||||
"POST", "/api/lm/recipe/{recipe_id}/opened", "record_recipe_open"
|
||||
),
|
||||
RouteDefinition("POST", "/api/lm/recipe/move", "move_recipe"),
|
||||
RouteDefinition("POST", "/api/lm/recipes/move-bulk", "move_recipes_bulk"),
|
||||
RouteDefinition("POST", "/api/lm/recipe/lora/reconnect", "reconnect_lora"),
|
||||
|
||||
@@ -8,11 +8,13 @@ import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Union, cast
|
||||
from ..config import config
|
||||
from ..utils.constants import VALID_CHECKPOINT_SUB_TYPES, VALID_LORA_TYPES
|
||||
from ..utils.file_utils import calculate_autov3
|
||||
from ..utils.recipe_open_stats import RecipeOpenStats
|
||||
from .recipe_cache import RecipeCache
|
||||
from .recipes.errors import RecipeNotFoundError, RecipePersistenceError
|
||||
from natsort import natsorted
|
||||
@@ -2781,7 +2783,11 @@ class RecipeScanner:
|
||||
Args:
|
||||
page: Current page number (1-based)
|
||||
page_size: Number of items per page
|
||||
sort_by: Sort method ('name' or 'date')
|
||||
sort_by: Sort method ('name', 'date', 'loras_count', 'opened',
|
||||
or 'random' with an optional seed like 'random:abc123'; the
|
||||
part after 'random:' is the shuffle seed, not a direction).
|
||||
'opened' hides recipes that were never opened — it is a
|
||||
"recently opened" view, not a plain reorder
|
||||
search: Search term
|
||||
filters: Dictionary of filters to apply
|
||||
search_options: Dictionary of search options to apply
|
||||
@@ -2962,7 +2968,7 @@ class RecipeScanner:
|
||||
]
|
||||
|
||||
# Apply sorting if not already handled by pre-sorted cache
|
||||
if ":" in sort_by or sort_field == "loras_count":
|
||||
if ":" in sort_by or sort_field in ("loras_count", "random", "opened"):
|
||||
field, order = (sort_by.split(":") + ["desc"])[:2]
|
||||
reverse = order.lower() == "desc"
|
||||
|
||||
@@ -2981,10 +2987,30 @@ class RecipeScanner:
|
||||
),
|
||||
reverse=reverse,
|
||||
)
|
||||
elif field == "opened":
|
||||
# "Recently Opened" view: recipes never opened are hidden.
|
||||
# The open stats live outside recipe metadata; see
|
||||
# RecipeOpenStats.
|
||||
opened_map = RecipeOpenStats().get_opened_map()
|
||||
filtered_data = [
|
||||
item
|
||||
for item in filtered_data
|
||||
if opened_map.get(str(item.get("id", ""))) is not None
|
||||
]
|
||||
filtered_data.sort(
|
||||
key=lambda x: opened_map.get(str(x.get("id", "")), 0),
|
||||
reverse=reverse,
|
||||
)
|
||||
elif field == "loras_count":
|
||||
filtered_data.sort(
|
||||
key=lambda x: len(x.get("loras", [])), reverse=reverse
|
||||
)
|
||||
elif field == "random":
|
||||
# Seeded random shuffle: same seed -> same order (stable
|
||||
# pagination across requests), matching the model pages.
|
||||
seed = order if order.lower() not in ("asc", "desc") else None
|
||||
rng = random.Random(seed or "random")
|
||||
rng.shuffle(filtered_data)
|
||||
|
||||
# Calculate pagination
|
||||
total_items = len(filtered_data)
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
"""Track recipe modal open timestamps for the "Recently Opened" sort.
|
||||
|
||||
The data is deliberately kept OUTSIDE the recipe metadata files: recording an
|
||||
open must be cheap and must never rewrite recipe JSON or EXIF (which the
|
||||
generic metadata update path does). A tiny JSON map of
|
||||
``recipe_id -> unix timestamp`` lives under
|
||||
``{settings_dir}/stats/recipe_last_opened.json`` and is written atomically on
|
||||
a short debounce.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
|
||||
from ..utils.settings_paths import get_settings_dir
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RecipeOpenStats:
|
||||
"""Persist the last time each recipe was opened in the recipe modal."""
|
||||
|
||||
STATS_FILENAME: str = "recipe_last_opened.json"
|
||||
SAVE_DELAY: float = 1.0 # seconds of debounce between consecutive writes
|
||||
|
||||
_instance: "RecipeOpenStats | None" = None
|
||||
_opened: dict[str, float]
|
||||
_file_mtime: float | None
|
||||
_dirty: bool
|
||||
_lock: asyncio.Lock
|
||||
_save_task: "asyncio.Task[None] | None"
|
||||
_stats_file_path: str
|
||||
_initialized: bool
|
||||
|
||||
def __new__(cls) -> "RecipeOpenStats":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._initialized = False
|
||||
return cls._instance
|
||||
|
||||
def __init__(self) -> None:
|
||||
if getattr(self, "_initialized", False):
|
||||
return
|
||||
self._opened = {}
|
||||
self._file_mtime = None
|
||||
self._dirty = False
|
||||
self._lock = asyncio.Lock()
|
||||
self._save_task = None
|
||||
self._stats_file_path = self._get_stats_file_path()
|
||||
self._load_stats()
|
||||
self._initialized = True
|
||||
|
||||
def _get_stats_file_path(self) -> str:
|
||||
settings_dir = get_settings_dir(create=True)
|
||||
return os.path.join(settings_dir, "stats", self.STATS_FILENAME)
|
||||
|
||||
def _load_stats(self) -> None:
|
||||
"""Load the opened map from disk, tolerating corrupt/absent files.
|
||||
|
||||
The mtime is recorded even when parsing fails so a corrupt file is
|
||||
not re-read (and re-logged) on every lookup.
|
||||
"""
|
||||
if not os.path.exists(self._stats_file_path):
|
||||
return
|
||||
try:
|
||||
mtime = os.path.getmtime(self._stats_file_path)
|
||||
except OSError:
|
||||
return
|
||||
try:
|
||||
with open(self._stats_file_path, "r", encoding="utf-8") as file_obj:
|
||||
raw = json.load(file_obj)
|
||||
if isinstance(raw, dict):
|
||||
self._opened = {
|
||||
str(key): float(value)
|
||||
for key, value in raw.items()
|
||||
if isinstance(value, (int, float))
|
||||
}
|
||||
except Exception as exc: # pragma: no cover - defensive logging path
|
||||
logger.error("Error loading recipe open stats: %s", exc)
|
||||
self._opened = {}
|
||||
self._file_mtime = mtime
|
||||
|
||||
def get_opened_map(self) -> dict[str, float]:
|
||||
"""Return a copy of ``recipe_id -> last opened timestamp``.
|
||||
|
||||
Refreshes from disk when the file changed since the last load so a
|
||||
second server process (or manual edit) is picked up without restart.
|
||||
"""
|
||||
try:
|
||||
if os.path.exists(self._stats_file_path):
|
||||
mtime = os.path.getmtime(self._stats_file_path)
|
||||
if self._file_mtime is None or mtime != self._file_mtime:
|
||||
self._load_stats()
|
||||
except OSError:
|
||||
pass
|
||||
return dict(self._opened)
|
||||
|
||||
def record_open(self, recipe_id: str) -> None:
|
||||
"""Mark a recipe as opened now; persists shortly in the background."""
|
||||
if not recipe_id:
|
||||
return
|
||||
self._opened[str(recipe_id)] = time.time()
|
||||
self._dirty = True
|
||||
if self._save_task is None or self._save_task.done():
|
||||
self._save_task = asyncio.create_task(self._delayed_save())
|
||||
|
||||
async def _delayed_save(self) -> None:
|
||||
"""Debounced writer: batches rapid consecutive opens into one write."""
|
||||
await asyncio.sleep(self.SAVE_DELAY)
|
||||
_ = await self.save_stats()
|
||||
|
||||
async def save_stats(self, force: bool = False) -> bool:
|
||||
"""Persist the opened map atomically if dirty (or when forced).
|
||||
|
||||
The on-disk map is merged in first so a second process sharing the
|
||||
settings dir does not lose its entries; the larger timestamp wins
|
||||
per recipe.
|
||||
"""
|
||||
if not force and not self._dirty:
|
||||
return False
|
||||
async with self._lock:
|
||||
if not force and not self._dirty:
|
||||
return False
|
||||
try:
|
||||
merged = self._merge_with_disk()
|
||||
os.makedirs(os.path.dirname(self._stats_file_path), exist_ok=True)
|
||||
temp_path = f"{self._stats_file_path}.tmp"
|
||||
with open(temp_path, "w", encoding="utf-8") as file_obj:
|
||||
json.dump(merged, file_obj, indent=2)
|
||||
os.replace(temp_path, self._stats_file_path)
|
||||
self._opened = merged
|
||||
self._file_mtime = os.path.getmtime(self._stats_file_path)
|
||||
self._dirty = False
|
||||
return True
|
||||
except Exception as exc: # pragma: no cover - defensive logging path
|
||||
logger.error("Error saving recipe open stats: %s", exc, exc_info=True)
|
||||
return False
|
||||
|
||||
def _merge_with_disk(self) -> dict[str, float]:
|
||||
"""Merge the in-memory map with the current on-disk map."""
|
||||
disk: dict[str, float] = {}
|
||||
try:
|
||||
if os.path.exists(self._stats_file_path):
|
||||
with open(self._stats_file_path, "r", encoding="utf-8") as file_obj:
|
||||
raw = json.load(file_obj)
|
||||
if isinstance(raw, dict):
|
||||
disk = {
|
||||
str(key): float(value)
|
||||
for key, value in raw.items()
|
||||
if isinstance(value, (int, float))
|
||||
}
|
||||
except Exception as exc: # pragma: no cover - defensive logging path
|
||||
logger.error("Error reading recipe open stats for merge: %s", exc)
|
||||
merged = dict(disk)
|
||||
for key, value in self._opened.items():
|
||||
merged[key] = max(value, disk.get(key, 0.0))
|
||||
return merged
|
||||
@@ -306,6 +306,14 @@ class RecipeModal {
|
||||
modalManager.showModal('recipeModal');
|
||||
|
||||
if (this.recipeId) {
|
||||
// Fire-and-forget: record this open for the "Recently Opened"
|
||||
// sort. Tracking must never disturb the modal, so failures are
|
||||
// swallowed.
|
||||
fetch(`/api/lm/recipe/${encodeURIComponent(this.recipeId)}/opened`, {
|
||||
method: 'POST',
|
||||
keepalive: true,
|
||||
}).catch(() => {});
|
||||
|
||||
const hydrationRequestId = ++this.recipeHydrationRequestId;
|
||||
const requestEditVersions = this.captureLocalEditVersions();
|
||||
this.hydrateRecipeDetails(
|
||||
|
||||
@@ -4,7 +4,7 @@ import { getStorageItem, setStorageItem, removeStorageItem, getSessionItem, setS
|
||||
import { showToast, openCivitaiByMetadata } from '../../utils/uiHelpers.js';
|
||||
import { performModelUpdateCheck } from '../../utils/updateCheckHelpers.js';
|
||||
import { sidebarManager } from '../SidebarManager.js';
|
||||
import { initSortDropdown } from './SortDropdown.js';
|
||||
import { initSortDropdown, applySortToSelect, randomizeSortValue } from './SortDropdown.js';
|
||||
|
||||
/**
|
||||
* PageControls class - Unified control management for model pages
|
||||
@@ -108,20 +108,20 @@ export class PageControls {
|
||||
const sortSelect = document.getElementById('sortSelect');
|
||||
if (sortSelect) {
|
||||
initSortDropdown(sortSelect);
|
||||
this.applySortToSelect(this.pageState.sortBy);
|
||||
applySortToSelect(this.pageState.sortBy);
|
||||
sortSelect.addEventListener('change', async (e) => {
|
||||
let value = e.target.value;
|
||||
if (value.startsWith('random')) {
|
||||
// Every pick of Random reshuffles the list: generate a
|
||||
// fresh seed so the backend keeps a stable order across
|
||||
// paginated requests.
|
||||
value = this._randomizeSortValue();
|
||||
value = randomizeSortValue();
|
||||
}
|
||||
this.pageState.sortBy = value;
|
||||
this.saveSortPreference(value);
|
||||
// Reset the seeded Random option when switching away from
|
||||
// Random, or re-apply the fresh seed when picking it again.
|
||||
this.applySortToSelect(value);
|
||||
applySortToSelect(value);
|
||||
await this.resetAndReload();
|
||||
});
|
||||
}
|
||||
@@ -322,44 +322,6 @@ export class PageControls {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Apply a sort value to the native sort <select>, keeping the Random
|
||||
* option's value in sync when the persisted value carries a seed
|
||||
* (e.g. "random:abc123"). Must be used instead of assigning
|
||||
* sortSelect.value directly whenever the value may be a seeded random
|
||||
* sort, otherwise the native select has no matching option.
|
||||
* @param {string} sortValue - Sort value like "name:asc" or "random:<seed>"
|
||||
*/
|
||||
applySortToSelect(sortValue) {
|
||||
const sortSelect = document.getElementById('sortSelect');
|
||||
if (!sortSelect) return;
|
||||
const randomOpt = sortSelect.querySelector('option[value="random"], option[value^="random:"]');
|
||||
if (randomOpt) {
|
||||
randomOpt.value = String(sortValue).startsWith('random') ? sortValue : 'random';
|
||||
}
|
||||
sortSelect.value = sortValue;
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a fresh seeded random sort value ("random:<seed>") and keep
|
||||
* the native <select> in sync so its value matches the persisted sort
|
||||
* string and the dropdown shows the selected label.
|
||||
* @returns {string} The new sort value, e.g. "random:abc123xyz"
|
||||
*/
|
||||
_randomizeSortValue() {
|
||||
const seed = Math.random().toString(36).slice(2, 12);
|
||||
const value = `random:${seed}`;
|
||||
const sortSelect = document.getElementById('sortSelect');
|
||||
if (sortSelect) {
|
||||
const randomOpt = sortSelect.querySelector('option[value="random"], option[value^="random:"]');
|
||||
if (randomOpt) {
|
||||
randomOpt.value = value;
|
||||
}
|
||||
sortSelect.value = value;
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
/**
|
||||
* Load sort preference from storage
|
||||
*/
|
||||
@@ -374,7 +336,7 @@ export class PageControls {
|
||||
// Handle legacy format conversion
|
||||
const convertedSort = this.convertLegacySortFormat(savedSort);
|
||||
this.pageState.sortBy = convertedSort;
|
||||
this.applySortToSelect(convertedSort);
|
||||
applySortToSelect(convertedSort);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -568,7 +530,7 @@ export class PageControls {
|
||||
this.pageState.sortBy = restoredSort;
|
||||
this.saveSortPreference(restoredSort);
|
||||
this._removeVlmSortOption();
|
||||
this.applySortToSelect(restoredSort);
|
||||
applySortToSelect(restoredSort);
|
||||
const sortSelect = document.getElementById('sortSelect');
|
||||
if (sortSelect) {
|
||||
sortSelect.disabled = false;
|
||||
@@ -620,7 +582,7 @@ export class PageControls {
|
||||
const savedGroupedSort = getStorageItem(groupedKey);
|
||||
if (savedGroupedSort) {
|
||||
this.pageState.sortBy = savedGroupedSort;
|
||||
this.applySortToSelect(savedGroupedSort);
|
||||
applySortToSelect(savedGroupedSort);
|
||||
}
|
||||
} else {
|
||||
// Leaving group mode: persist current sort for next time, restore non-group sort
|
||||
@@ -628,7 +590,7 @@ export class PageControls {
|
||||
const savedNormalSort = getStorageItem(`${this.pageType}_sort`);
|
||||
if (savedNormalSort) {
|
||||
this.pageState.sortBy = savedNormalSort;
|
||||
this.applySortToSelect(savedNormalSort);
|
||||
applySortToSelect(savedNormalSort);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -913,7 +875,7 @@ export class PageControls {
|
||||
}
|
||||
|
||||
if (sortSelect) {
|
||||
this.applySortToSelect(this.pageState.sortBy);
|
||||
applySortToSelect(this.pageState.sortBy);
|
||||
}
|
||||
if (searchInput) {
|
||||
searchInput.value = this.pageState.filters?.search || '';
|
||||
|
||||
@@ -18,6 +18,44 @@
|
||||
const SORT_GROUP_SELECTOR = '.sort-dropdown-group';
|
||||
const ACTIVE_GROUP_SELECTOR = '.sort-dropdown-group.active, .dropdown-group.active';
|
||||
|
||||
/**
|
||||
* Apply a sort value to the page's native sort <select>, keeping the Random
|
||||
* option's value in sync when the persisted value carries a seed
|
||||
* (e.g. "random:abc123"). Must be used instead of assigning
|
||||
* sortSelect.value directly whenever the value may be a seeded random
|
||||
* sort, otherwise the native select has no matching option.
|
||||
* @param {string} sortValue - Sort value like "name:asc" or "random:<seed>"
|
||||
*/
|
||||
export function applySortToSelect(sortValue) {
|
||||
const sortSelect = document.getElementById('sortSelect');
|
||||
if (!sortSelect) return;
|
||||
const randomOpt = sortSelect.querySelector('option[value="random"], option[value^="random:"]');
|
||||
if (randomOpt) {
|
||||
randomOpt.value = String(sortValue).startsWith('random') ? sortValue : 'random';
|
||||
}
|
||||
sortSelect.value = sortValue;
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a fresh seeded random sort value ("random:<seed>") and keep the
|
||||
* native <select> in sync so its value matches the persisted sort string and
|
||||
* the dropdown shows the selected label.
|
||||
* @returns {string} The new sort value, e.g. "random:abc123xyz"
|
||||
*/
|
||||
export function randomizeSortValue() {
|
||||
const seed = Math.random().toString(36).slice(2, 12);
|
||||
const value = `random:${seed}`;
|
||||
const sortSelect = document.getElementById('sortSelect');
|
||||
if (sortSelect) {
|
||||
const randomOpt = sortSelect.querySelector('option[value="random"], option[value^="random:"]');
|
||||
if (randomOpt) {
|
||||
randomOpt.value = value;
|
||||
}
|
||||
sortSelect.value = value;
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
/**
|
||||
* Initialize a decoupled sort dropdown around a native <select>.
|
||||
* Idempotent: safe to call more than once on the same element.
|
||||
|
||||
+14
-4
@@ -10,7 +10,7 @@ import { DuplicatesManager } from './components/DuplicatesManager.js';
|
||||
import { refreshVirtualScroll, recreateVirtualScroll } from './utils/infiniteScroll.js';
|
||||
import { refreshRecipes, RecipeSidebarApiClient } from './api/recipeApi.js';
|
||||
import { sidebarManager } from './components/SidebarManager.js';
|
||||
import { initSortDropdown } from './components/controls/SortDropdown.js';
|
||||
import { initSortDropdown, applySortToSelect, randomizeSortValue } from './components/controls/SortDropdown.js';
|
||||
|
||||
class RecipePageControls {
|
||||
constructor() {
|
||||
@@ -245,10 +245,20 @@ class RecipeManager {
|
||||
this.pageState.sortBy = savedSort;
|
||||
}
|
||||
initSortDropdown(sortSelect);
|
||||
sortSelect.value = this.pageState.sortBy || 'date:desc';
|
||||
applySortToSelect(this.pageState.sortBy || 'date:desc');
|
||||
sortSelect.addEventListener('change', () => {
|
||||
this.pageState.sortBy = sortSelect.value;
|
||||
setStorageItem('recipes_sort', sortSelect.value);
|
||||
let value = sortSelect.value;
|
||||
if (value.startsWith('random')) {
|
||||
// Every pick of Random reshuffles the list: generate a
|
||||
// fresh seed so the backend keeps a stable order across
|
||||
// paginated requests.
|
||||
value = randomizeSortValue();
|
||||
}
|
||||
this.pageState.sortBy = value;
|
||||
setStorageItem('recipes_sort', value);
|
||||
// Reset the seeded Random option when switching away from
|
||||
// Random, or re-apply the fresh seed when picking it again.
|
||||
applySortToSelect(value);
|
||||
refreshVirtualScroll();
|
||||
});
|
||||
}
|
||||
|
||||
@@ -646,10 +646,17 @@ export class MasonryScroller {
|
||||
const pageType = state.currentPageType;
|
||||
|
||||
if (pageType === 'recipes') {
|
||||
placeholderText = `
|
||||
<p>No recipes found</p>
|
||||
<p>Add recipe images to your recipes folder to see them here.</p>
|
||||
`;
|
||||
if (String(getCurrentPageState().sortBy).startsWith('opened')) {
|
||||
placeholderText = `
|
||||
<p>No recently opened recipes</p>
|
||||
<p>Recipes you open will appear here.</p>
|
||||
`;
|
||||
} else {
|
||||
placeholderText = `
|
||||
<p>No recipes found</p>
|
||||
<p>Add recipe images to your recipes folder to see them here.</p>
|
||||
`;
|
||||
}
|
||||
} else if (pageType === 'loras') {
|
||||
placeholderText = `
|
||||
<p>No LoRAs found</p>
|
||||
|
||||
@@ -699,10 +699,17 @@ export class VirtualScroller {
|
||||
const pageType = state.currentPageType;
|
||||
|
||||
if (pageType === 'recipes') {
|
||||
placeholderText = `
|
||||
<p>No recipes found</p>
|
||||
<p>Add recipe images to your recipes folder to see them here.</p>
|
||||
`;
|
||||
if (String(getCurrentPageState().sortBy).startsWith('opened')) {
|
||||
placeholderText = `
|
||||
<p>No recently opened recipes</p>
|
||||
<p>Recipes you open will appear here.</p>
|
||||
`;
|
||||
} else {
|
||||
placeholderText = `
|
||||
<p>No recipes found</p>
|
||||
<p>Add recipe images to your recipes folder to see them here.</p>
|
||||
`;
|
||||
}
|
||||
} else if (pageType === 'loras') {
|
||||
placeholderText = `
|
||||
<p>No LoRAs found</p>
|
||||
|
||||
@@ -48,17 +48,20 @@
|
||||
<option value="versions_count:asc">{{ t('loras.controls.sort.versionsCountAsc', default='Fewest versions first') }}</option>
|
||||
</optgroup>
|
||||
{% endif %}
|
||||
{% if page_id != 'recipes' %}
|
||||
<optgroup label="{{ t('loras.controls.sort.random', default='Random') }}">
|
||||
<option value="random">{{ t('loras.controls.sort.randomAction', default='Randomize (shuffle)') }}</option>
|
||||
</optgroup>
|
||||
{% endif %}
|
||||
{% if page_id == 'recipes' %}
|
||||
<optgroup label="{{ t('recipes.controls.sort.lorasCount') }}">
|
||||
<option value="loras_count:desc">{{ t('recipes.controls.sort.lorasCountDesc') }}</option>
|
||||
<option value="loras_count:asc">{{ t('recipes.controls.sort.lorasCountAsc') }}</option>
|
||||
</optgroup>
|
||||
{% endif %}
|
||||
{% if page_id == 'recipes' %}
|
||||
<optgroup label="{{ t('recipes.controls.sort.opened', default='Recently Opened') }}">
|
||||
<option value="opened:desc">{{ t('recipes.controls.sort.openedDesc', default='Recently opened') }}</option>
|
||||
</optgroup>
|
||||
{% endif %}
|
||||
<optgroup label="{{ t('loras.controls.sort.random', default='Random') }}">
|
||||
<option value="random">{{ t('loras.controls.sort.randomAction', default='Randomize (shuffle)') }}</option>
|
||||
</optgroup>
|
||||
</select>
|
||||
</div>
|
||||
<div title="{% if page_id == 'recipes' %}{{ t('recipes.controls.refresh.title') }}{% else %}{{ t('loras.controls.refresh.title') }}{% endif %}" class="control-group dropdown-group">
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { describe, it, beforeEach, afterEach, expect, vi } from 'vitest';
|
||||
import { applySortToSelect } from '../../../static/js/components/controls/SortDropdown.js';
|
||||
|
||||
const resetAndReloadMock = vi.fn();
|
||||
const getModelApiClientMock = vi.fn();
|
||||
@@ -190,7 +191,7 @@ describe('Random sort option', () => {
|
||||
sortSelect.value = 'random';
|
||||
sortSelect.dispatchEvent(new Event('change', { bubbles: true }));
|
||||
await Promise.resolve();
|
||||
controls.applySortToSelect('name:desc');
|
||||
applySortToSelect('name:desc');
|
||||
|
||||
expect(sortSelect.value).toBe('name:desc');
|
||||
expect(randomOpt.value).toBe('random');
|
||||
|
||||
@@ -0,0 +1,236 @@
|
||||
import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
|
||||
import { renderRecipesPage } from '../utils/pageFixtures.js';
|
||||
import { applySortToSelect } from '../../../static/js/components/controls/SortDropdown.js';
|
||||
|
||||
const initializeAppMock = vi.fn();
|
||||
const initializePageFeaturesMock = vi.fn();
|
||||
const getCurrentPageStateMock = vi.fn();
|
||||
const getSessionItemMock = vi.fn();
|
||||
const removeSessionItemMock = vi.fn();
|
||||
const getStorageItemMock = vi.fn();
|
||||
const setStorageItemMock = vi.fn();
|
||||
const removeStorageItemMock = vi.fn();
|
||||
const refreshVirtualScrollMock = vi.fn();
|
||||
const refreshRecipesMock = vi.fn();
|
||||
|
||||
let importManagerInstance;
|
||||
let recipeModalInstance;
|
||||
let duplicatesManagerInstance;
|
||||
|
||||
const ImportManagerMock = vi.fn(() => importManagerInstance);
|
||||
const RecipeModalMock = vi.fn(() => recipeModalInstance);
|
||||
const DuplicatesManagerMock = vi.fn(() => duplicatesManagerInstance);
|
||||
|
||||
vi.mock('../../../static/js/core.js', () => ({
|
||||
appCore: {
|
||||
initialize: initializeAppMock,
|
||||
initializePageFeatures: initializePageFeaturesMock,
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/managers/ImportManager.js', () => ({
|
||||
ImportManager: ImportManagerMock,
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/components/RecipeModal.js', () => ({
|
||||
RecipeModal: RecipeModalMock,
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/state/index.js', () => ({
|
||||
getCurrentPageState: getCurrentPageStateMock,
|
||||
state: {
|
||||
currentPageType: 'recipes',
|
||||
global: { settings: {} },
|
||||
virtualScroller: {
|
||||
removeItemByFilePath: vi.fn(),
|
||||
updateSingleItem: vi.fn(),
|
||||
refreshWithData: vi.fn(),
|
||||
},
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/utils/storageHelpers.js', () => ({
|
||||
getSessionItem: getSessionItemMock,
|
||||
removeSessionItem: removeSessionItemMock,
|
||||
getStorageItem: getStorageItemMock,
|
||||
setStorageItem: setStorageItemMock,
|
||||
removeStorageItem: removeStorageItemMock,
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/components/ContextMenu/index.js', () => ({
|
||||
RecipeContextMenu: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/components/DuplicatesManager.js', () => ({
|
||||
DuplicatesManager: DuplicatesManagerMock,
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/utils/infiniteScroll.js', () => ({
|
||||
refreshVirtualScroll: refreshVirtualScrollMock,
|
||||
recreateVirtualScroll: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/api/recipeApi.js', () => ({
|
||||
refreshRecipes: refreshRecipesMock,
|
||||
RecipeSidebarApiClient: vi.fn(() => ({
|
||||
apiConfig: { config: { displayName: 'Recipes', supportsMove: true } },
|
||||
fetchUnifiedFolderTree: vi.fn().mockResolvedValue({ success: true, tree: {} }),
|
||||
fetchModelFolders: vi.fn().mockResolvedValue({ success: true, folders: [] }),
|
||||
fetchModelRoots: vi.fn().mockResolvedValue({ roots: ['/recipes'] }),
|
||||
moveBulkModels: vi.fn(),
|
||||
moveSingleModel: vi.fn(),
|
||||
})),
|
||||
}));
|
||||
|
||||
vi.mock('../../../static/js/components/SidebarManager.js', () => ({
|
||||
sidebarManager: {
|
||||
setHostPageControls: vi.fn(),
|
||||
initialize: vi.fn(async () => {}),
|
||||
refresh: vi.fn(async () => {}),
|
||||
cleanup: vi.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
function renderSortSelect() {
|
||||
const sortSelectElement = document.createElement('select');
|
||||
sortSelectElement.id = 'sortSelect';
|
||||
sortSelectElement.innerHTML = `
|
||||
<option value="date:desc">Newest</option>
|
||||
<option value="name:asc">Name A-Z</option>
|
||||
<option value="random">Randomize (shuffle)</option>
|
||||
`;
|
||||
document.body.appendChild(sortSelectElement);
|
||||
return sortSelectElement;
|
||||
}
|
||||
|
||||
describe('RecipeManager Random sort', () => {
|
||||
let RecipeManager;
|
||||
let pageState;
|
||||
|
||||
beforeEach(async () => {
|
||||
vi.resetModules();
|
||||
vi.clearAllMocks();
|
||||
|
||||
importManagerInstance = { showImportModal: vi.fn() };
|
||||
recipeModalInstance = { showRecipeDetails: vi.fn() };
|
||||
duplicatesManagerInstance = {
|
||||
findDuplicates: vi.fn(),
|
||||
selectLatestDuplicates: vi.fn(),
|
||||
deleteSelectedDuplicates: vi.fn(),
|
||||
confirmDeleteDuplicates: vi.fn(),
|
||||
exitDuplicateMode: vi.fn(),
|
||||
};
|
||||
|
||||
pageState = {
|
||||
sortBy: 'date:desc',
|
||||
searchOptions: undefined,
|
||||
customFilter: undefined,
|
||||
duplicatesMode: false,
|
||||
};
|
||||
|
||||
getCurrentPageStateMock.mockImplementation(() => pageState);
|
||||
initializeAppMock.mockResolvedValue(undefined);
|
||||
initializePageFeaturesMock.mockResolvedValue(undefined);
|
||||
refreshVirtualScrollMock.mockImplementation(() => {});
|
||||
refreshRecipesMock.mockResolvedValue('refreshed');
|
||||
getSessionItemMock.mockImplementation(() => null);
|
||||
removeSessionItemMock.mockImplementation(() => {});
|
||||
getStorageItemMock.mockImplementation(() => null);
|
||||
setStorageItemMock.mockImplementation(() => {});
|
||||
|
||||
renderRecipesPage();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
document.body.innerHTML = '';
|
||||
delete window.recipeManager;
|
||||
delete window.importManager;
|
||||
});
|
||||
|
||||
async function createManager() {
|
||||
({ RecipeManager } = await import('../../../static/js/recipes.js'));
|
||||
const manager = new RecipeManager();
|
||||
await manager.initialize();
|
||||
return manager;
|
||||
}
|
||||
|
||||
it('generates a seeded sort value when Random is picked', async () => {
|
||||
const sortSelect = renderSortSelect();
|
||||
const randomOpt = sortSelect.querySelector('option[value="random"]');
|
||||
await createManager();
|
||||
|
||||
sortSelect.value = 'random';
|
||||
sortSelect.dispatchEvent(new Event('change', { bubbles: true }));
|
||||
await Promise.resolve();
|
||||
|
||||
expect(pageState.sortBy).toMatch(/^random:[a-z0-9]+$/);
|
||||
expect(setStorageItemMock).toHaveBeenCalledWith('recipes_sort', pageState.sortBy);
|
||||
expect(randomOpt.value).toBe(pageState.sortBy);
|
||||
expect(sortSelect.value).toBe(pageState.sortBy);
|
||||
expect(refreshVirtualScrollMock).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('reshuffles with a fresh seed every time Random is picked again', async () => {
|
||||
const sortSelect = renderSortSelect();
|
||||
const randomOpt = sortSelect.querySelector('option[value="random"]');
|
||||
await createManager();
|
||||
|
||||
sortSelect.value = 'random';
|
||||
sortSelect.dispatchEvent(new Event('change', { bubbles: true }));
|
||||
await Promise.resolve();
|
||||
const firstSeed = pageState.sortBy;
|
||||
|
||||
sortSelect.value = randomOpt.value;
|
||||
sortSelect.dispatchEvent(new Event('change', { bubbles: true }));
|
||||
await Promise.resolve();
|
||||
|
||||
expect(pageState.sortBy).toMatch(/^random:[a-z0-9]+$/);
|
||||
expect(pageState.sortBy).not.toBe(firstSeed);
|
||||
});
|
||||
|
||||
it('restores a persisted seeded random sort on load', async () => {
|
||||
const sortSelect = renderSortSelect();
|
||||
const savedSort = 'random:persistedseed';
|
||||
getStorageItemMock.mockImplementation((key) =>
|
||||
key === 'recipes_sort' ? savedSort : null
|
||||
);
|
||||
await createManager();
|
||||
|
||||
expect(pageState.sortBy).toBe(savedSort);
|
||||
expect(sortSelect.value).toBe(savedSort);
|
||||
expect(sortSelect.querySelector('option[value="random:persistedseed"]')).not.toBeNull();
|
||||
});
|
||||
|
||||
it('applies a non-random sort back to the plain random option', async () => {
|
||||
const sortSelect = renderSortSelect();
|
||||
const randomOpt = sortSelect.querySelector('option[value="random"]');
|
||||
await createManager();
|
||||
|
||||
sortSelect.value = 'random';
|
||||
sortSelect.dispatchEvent(new Event('change', { bubbles: true }));
|
||||
await Promise.resolve();
|
||||
applySortToSelect('name:asc');
|
||||
|
||||
expect(sortSelect.value).toBe('name:asc');
|
||||
expect(randomOpt.value).toBe('random');
|
||||
});
|
||||
|
||||
it('resets the seeded option when switching away from Random via the change handler', async () => {
|
||||
const sortSelect = renderSortSelect();
|
||||
const randomOpt = sortSelect.querySelector('option[value="random"]');
|
||||
await createManager();
|
||||
|
||||
sortSelect.value = 'random';
|
||||
sortSelect.dispatchEvent(new Event('change', { bubbles: true }));
|
||||
await Promise.resolve();
|
||||
expect(randomOpt.value).toMatch(/^random:[a-z0-9]+$/);
|
||||
|
||||
sortSelect.value = 'name:asc';
|
||||
sortSelect.dispatchEvent(new Event('change', { bubbles: true }));
|
||||
await Promise.resolve();
|
||||
|
||||
expect(pageState.sortBy).toBe('name:asc');
|
||||
expect(sortSelect.value).toBe('name:asc');
|
||||
expect(randomOpt.value).toBe('random');
|
||||
});
|
||||
});
|
||||
@@ -324,6 +324,18 @@ describe('MasonryScroller', () => {
|
||||
expect(placeholder.textContent).toContain('No recipes found');
|
||||
});
|
||||
|
||||
it('shows the recently-opened empty placeholder under the opened sort', async () => {
|
||||
getCurrentPageState().sortBy = 'opened:desc';
|
||||
const { scroller, grid } = track(createScroller({ items: [] }));
|
||||
|
||||
await scroller.initialize();
|
||||
|
||||
const placeholder = grid.querySelector('#virtualScrollPlaceholder');
|
||||
expect(placeholder).not.toBeNull();
|
||||
expect(placeholder.textContent).toContain('No recently opened recipes');
|
||||
getCurrentPageState().sortBy = '';
|
||||
});
|
||||
|
||||
it('dispose removes classes, spacer and event listeners', () => {
|
||||
const { scroller, grid } = track(createScroller());
|
||||
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
"""Tests for the Random Checkpoint/Unet Loader nodes' base-model filtering and
|
||||
random-selection behavior.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from py.nodes.random_checkpoint_loader import RandomCheckpointLoaderLM
|
||||
from py.nodes.random_unet_loader import RandomUNETLoaderLM
|
||||
|
||||
|
||||
class _FakeCache:
|
||||
def __init__(self, raw_data):
|
||||
self.raw_data = raw_data
|
||||
|
||||
|
||||
class _FakeScanner:
|
||||
def __init__(self, raw_data, model_roots):
|
||||
self._raw_data = raw_data
|
||||
self._model_roots = model_roots
|
||||
|
||||
async def get_cached_data(self, force_refresh=False):
|
||||
return _FakeCache(self._raw_data)
|
||||
|
||||
def get_model_roots(self):
|
||||
return self._model_roots
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def base_model_library(tmp_path, monkeypatch):
|
||||
from py.services.service_registry import ServiceRegistry
|
||||
|
||||
illustrious = tmp_path / "illustrious.safetensors"
|
||||
illustrious.write_bytes(b"x")
|
||||
flux = tmp_path / "flux.safetensors"
|
||||
flux.write_bytes(b"x")
|
||||
missing = tmp_path / "missing.safetensors" # referenced but never created
|
||||
|
||||
raw_data = [
|
||||
{
|
||||
"sub_type": "checkpoint",
|
||||
"file_path": str(illustrious),
|
||||
"base_model": "Illustrious",
|
||||
},
|
||||
{"sub_type": "checkpoint", "file_path": str(flux), "base_model": "Flux.1 D"},
|
||||
{
|
||||
"sub_type": "checkpoint",
|
||||
"file_path": str(missing),
|
||||
"base_model": "SDXL 1.0",
|
||||
},
|
||||
{
|
||||
"sub_type": "diffusion_model",
|
||||
"file_path": str(flux),
|
||||
"base_model": "Flux.1 D",
|
||||
},
|
||||
]
|
||||
|
||||
async def _fake_scanner():
|
||||
return _FakeScanner(raw_data, [str(tmp_path)])
|
||||
|
||||
monkeypatch.setattr(ServiceRegistry, "get_checkpoint_scanner", _fake_scanner)
|
||||
return tmp_path
|
||||
|
||||
|
||||
def test_checkpoint_names_drop_deleted_files(tmp_path, monkeypatch):
|
||||
from py.services.service_registry import ServiceRegistry
|
||||
|
||||
existing = tmp_path / "keep.safetensors"
|
||||
existing.write_bytes(b"x")
|
||||
deleted = tmp_path / "deleted.safetensors" # referenced but never created
|
||||
|
||||
raw_data = [
|
||||
{"sub_type": "checkpoint", "file_path": str(existing)},
|
||||
{"sub_type": "checkpoint", "file_path": str(deleted)},
|
||||
# Wrong type must stay excluded by the sub_type filter.
|
||||
{"sub_type": "diffusion_model", "file_path": str(existing)},
|
||||
]
|
||||
|
||||
async def _fake_scanner():
|
||||
return _FakeScanner(raw_data, [str(tmp_path)])
|
||||
|
||||
monkeypatch.setattr(ServiceRegistry, "get_checkpoint_scanner", _fake_scanner)
|
||||
assert RandomCheckpointLoaderLM._get_checkpoint_names() == ["keep.safetensors"]
|
||||
|
||||
|
||||
def test_unet_names_drop_deleted_files(tmp_path, monkeypatch):
|
||||
from py.services.service_registry import ServiceRegistry
|
||||
|
||||
existing = tmp_path / "keep.safetensors"
|
||||
existing.write_bytes(b"x")
|
||||
deleted = tmp_path / "deleted.safetensors"
|
||||
|
||||
raw_data = [
|
||||
{"sub_type": "diffusion_model", "file_path": str(existing)},
|
||||
{"sub_type": "diffusion_model", "file_path": str(deleted)},
|
||||
{"sub_type": "checkpoint", "file_path": str(existing)},
|
||||
]
|
||||
|
||||
async def _fake_scanner():
|
||||
return _FakeScanner(raw_data, [str(tmp_path)])
|
||||
|
||||
monkeypatch.setattr(ServiceRegistry, "get_checkpoint_scanner", _fake_scanner)
|
||||
assert RandomUNETLoaderLM._get_unet_names() == ["keep.safetensors"]
|
||||
|
||||
|
||||
def test_checkpoint_names_empty_when_scanner_fails(tmp_path, monkeypatch):
|
||||
from py.services.service_registry import ServiceRegistry
|
||||
|
||||
def _boom():
|
||||
raise RuntimeError("scanner not available")
|
||||
|
||||
monkeypatch.setattr(ServiceRegistry, "get_checkpoint_scanner", _boom)
|
||||
assert RandomCheckpointLoaderLM._get_checkpoint_names() == []
|
||||
|
||||
|
||||
def test_checkpoint_available_base_models(base_model_library):
|
||||
# "SDXL 1.0" is excluded because its file no longer exists on disk.
|
||||
assert RandomCheckpointLoaderLM._get_available_base_models() == [
|
||||
"Any",
|
||||
"Flux.1 D",
|
||||
"Illustrious",
|
||||
]
|
||||
|
||||
|
||||
def test_checkpoint_names_filtered_by_base_model(base_model_library):
|
||||
assert RandomCheckpointLoaderLM._get_checkpoint_names("Illustrious") == [
|
||||
"illustrious.safetensors"
|
||||
]
|
||||
assert RandomCheckpointLoaderLM._get_checkpoint_names("Any") == [
|
||||
"flux.safetensors",
|
||||
"illustrious.safetensors",
|
||||
]
|
||||
|
||||
|
||||
def test_unet_available_base_models(base_model_library):
|
||||
assert RandomUNETLoaderLM._get_available_base_models() == ["Any", "Flux.1 D"]
|
||||
|
||||
|
||||
def test_load_checkpoint_random_selection_uses_pool(base_model_library, monkeypatch):
|
||||
from py.nodes import random_checkpoint_loader as random_checkpoint_loader_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
random_checkpoint_loader_module,
|
||||
"get_checkpoint_info_absolute",
|
||||
lambda name: (str(base_model_library / name), {"file_path": name}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
random_checkpoint_loader_module.comfy.sd,
|
||||
"load_checkpoint_guess_config",
|
||||
lambda *a, **k: ("MODEL", "CLIP", "VAE", None),
|
||||
raising=False,
|
||||
)
|
||||
|
||||
node = RandomCheckpointLoaderLM()
|
||||
result = node.load_checkpoint(
|
||||
"ignored.safetensors", select_at_random=True, base_model="Illustrious"
|
||||
)
|
||||
# Only one checkpoint matches "Illustrious", so the random pick is deterministic here.
|
||||
assert result[3] == "illustrious.safetensors"
|
||||
|
||||
|
||||
def test_load_checkpoint_random_selection_raises_when_pool_empty(base_model_library):
|
||||
node = RandomCheckpointLoaderLM()
|
||||
with pytest.raises(FileNotFoundError, match="No checkpoints found"):
|
||||
node.load_checkpoint(
|
||||
"ignored.safetensors", select_at_random=True, base_model="SDXL 1.0"
|
||||
)
|
||||
|
||||
|
||||
def test_checkpoint_is_changed_forces_rerun_when_random():
|
||||
assert RandomCheckpointLoaderLM.IS_CHANGED(
|
||||
"a.safetensors", select_at_random=True, base_model="Any"
|
||||
) != RandomCheckpointLoaderLM.IS_CHANGED(
|
||||
"a.safetensors", select_at_random=True, base_model="Any"
|
||||
)
|
||||
assert RandomCheckpointLoaderLM.IS_CHANGED(
|
||||
"a.safetensors", select_at_random=False, base_model="Any"
|
||||
) == RandomCheckpointLoaderLM.IS_CHANGED(
|
||||
"a.safetensors", select_at_random=False, base_model="Any"
|
||||
)
|
||||
@@ -1047,6 +1047,106 @@ async def test_get_paginated_data_sorting(recipe_scanner):
|
||||
assert [i["id"] for i in res["items"]] == ["C", "A", "B"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_paginated_data_random_sort(recipe_scanner):
|
||||
scanner, _ = recipe_scanner
|
||||
|
||||
# Add test recipes
|
||||
for rid, title in [("A", "Alpha"), ("B", "Beta"), ("C", "Gamma")]:
|
||||
await scanner.add_recipe(
|
||||
{
|
||||
"id": rid,
|
||||
"title": title,
|
||||
"created_date": 10.0,
|
||||
"loras": [{}],
|
||||
"file_path": f"{rid.lower()}.png",
|
||||
}
|
||||
)
|
||||
|
||||
await asyncio.sleep(0)
|
||||
await _wait_for_resort(scanner)
|
||||
|
||||
# Same seed -> same order (deterministic, stable pagination)
|
||||
res1 = await scanner.get_paginated_data(
|
||||
page=1, page_size=10, sort_by="random:seed123"
|
||||
)
|
||||
res2 = await scanner.get_paginated_data(
|
||||
page=1, page_size=10, sort_by="random:seed123"
|
||||
)
|
||||
ids1 = [i["id"] for i in res1["items"]]
|
||||
ids2 = [i["id"] for i in res2["items"]]
|
||||
assert ids1 == ids2
|
||||
assert sorted(ids1) == ["A", "B", "C"]
|
||||
|
||||
# Plain "random" (no seed) also returns the full set
|
||||
res3 = await scanner.get_paginated_data(page=1, page_size=10, sort_by="random")
|
||||
assert sorted(i["id"] for i in res3["items"]) == ["A", "B", "C"]
|
||||
|
||||
# Stable pagination: page1 + page2 with the same seed concatenate to the
|
||||
# full seeded order, with no duplicates across pages
|
||||
p1 = await scanner.get_paginated_data(
|
||||
page=1, page_size=2, sort_by="random:seed123"
|
||||
)
|
||||
p2 = await scanner.get_paginated_data(
|
||||
page=2, page_size=2, sort_by="random:seed123"
|
||||
)
|
||||
combined = [i["id"] for i in p1["items"]] + [i["id"] for i in p2["items"]]
|
||||
assert combined == ids1
|
||||
assert len(set(combined)) == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_paginated_data_opened_sort(recipe_scanner, monkeypatch):
|
||||
scanner, _ = recipe_scanner
|
||||
|
||||
for rid, title in [("A", "Alpha"), ("B", "Beta"), ("C", "Gamma")]:
|
||||
await scanner.add_recipe(
|
||||
{
|
||||
"id": rid,
|
||||
"title": title,
|
||||
"created_date": 10.0,
|
||||
"loras": [{}],
|
||||
"file_path": f"{rid.lower()}.png",
|
||||
}
|
||||
)
|
||||
|
||||
await asyncio.sleep(0)
|
||||
await _wait_for_resort(scanner)
|
||||
|
||||
class _FakeStats:
|
||||
def get_opened_map(self):
|
||||
return {"B": 300.0, "C": 200.0}
|
||||
|
||||
monkeypatch.setattr(
|
||||
"py.services.recipe_scanner.RecipeOpenStats", lambda: _FakeStats()
|
||||
)
|
||||
|
||||
# Never-opened A is hidden from the view; B (300) > C (200)
|
||||
res = await scanner.get_paginated_data(page=1, page_size=10, sort_by="opened:desc")
|
||||
assert [i["id"] for i in res["items"]] == ["B", "C"]
|
||||
assert res["total"] == 2
|
||||
|
||||
# ASC: C (200) < B (300)
|
||||
res = await scanner.get_paginated_data(page=1, page_size=10, sort_by="opened:asc")
|
||||
assert [i["id"] for i in res["items"]] == ["C", "B"]
|
||||
|
||||
# Plain "opened" (no direction) behaves like desc by default
|
||||
res = await scanner.get_paginated_data(page=1, page_size=10, sort_by="opened")
|
||||
assert [i["id"] for i in res["items"]] == ["B", "C"]
|
||||
|
||||
# When nothing was opened the view is empty (not a fallback reorder)
|
||||
class _EmptyStats:
|
||||
def get_opened_map(self):
|
||||
return {}
|
||||
|
||||
monkeypatch.setattr(
|
||||
"py.services.recipe_scanner.RecipeOpenStats", lambda: _EmptyStats()
|
||||
)
|
||||
res = await scanner.get_paginated_data(page=1, page_size=10, sort_by="opened:desc")
|
||||
assert res["items"] == []
|
||||
assert res["total"] == 0
|
||||
|
||||
|
||||
async def test_build_image_id_map_filters_correctly(recipe_scanner):
|
||||
"""Only recipes with valid CivitAI source_path appear in image_id_map.
|
||||
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from py.utils import recipe_open_stats as stats_module
|
||||
from py.utils.recipe_open_stats import RecipeOpenStats
|
||||
|
||||
|
||||
async def _finalize(tasks) -> None:
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await task
|
||||
RecipeOpenStats._instance = None
|
||||
|
||||
|
||||
def _prepare(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
RecipeOpenStats._instance = None
|
||||
settings_dir = tmp_path / "settings"
|
||||
settings_dir.mkdir(parents=True, exist_ok=True)
|
||||
monkeypatch.setattr(
|
||||
stats_module, "get_settings_dir", lambda create=True: str(settings_dir)
|
||||
)
|
||||
created_tasks = []
|
||||
real_create_task = stats_module.asyncio.create_task
|
||||
|
||||
def _track_task(coro):
|
||||
task = real_create_task(coro)
|
||||
created_tasks.append(task)
|
||||
return task
|
||||
|
||||
monkeypatch.setattr(stats_module.asyncio, "create_task", _track_task)
|
||||
return RecipeOpenStats(), created_tasks, settings_dir
|
||||
|
||||
|
||||
async def _wait_for_save(stats_file: Path) -> None:
|
||||
for _ in range(100):
|
||||
if stats_file.exists():
|
||||
return
|
||||
await asyncio.sleep(0.01)
|
||||
raise AssertionError("Recipe open stats file was never written")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_open_persists_timestamp(tmp_path, monkeypatch):
|
||||
stats, tasks, settings_dir = _prepare(tmp_path, monkeypatch)
|
||||
stats_file = settings_dir / "stats" / RecipeOpenStats.STATS_FILENAME
|
||||
|
||||
stats.record_open("abc-123")
|
||||
await _wait_for_save(stats_file)
|
||||
|
||||
data = json.loads(stats_file.read_text(encoding="utf-8"))
|
||||
assert isinstance(data["abc-123"], float)
|
||||
await _finalize(tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_open_updates_existing_entry(tmp_path, monkeypatch):
|
||||
stats, tasks, settings_dir = _prepare(tmp_path, monkeypatch)
|
||||
stats_file = settings_dir / "stats" / RecipeOpenStats.STATS_FILENAME
|
||||
|
||||
stats.record_open("r1")
|
||||
await _wait_for_save(stats_file)
|
||||
first = json.loads(stats_file.read_text(encoding="utf-8"))["r1"]
|
||||
|
||||
await asyncio.sleep(0.01)
|
||||
stats.record_open("r1")
|
||||
await stats.save_stats(force=True)
|
||||
|
||||
second = json.loads(stats_file.read_text(encoding="utf-8"))["r1"]
|
||||
assert second > first
|
||||
await _finalize(tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_opened_map_reloads_on_file_change(tmp_path, monkeypatch):
|
||||
stats, tasks, settings_dir = _prepare(tmp_path, monkeypatch)
|
||||
stats_file = settings_dir / "stats" / RecipeOpenStats.STATS_FILENAME
|
||||
|
||||
stats.record_open("r1")
|
||||
await _wait_for_save(stats_file)
|
||||
|
||||
stats_file.write_text(json.dumps({"r2": 500.0}), encoding="utf-8")
|
||||
opened_map = stats.get_opened_map()
|
||||
assert opened_map == {"r2": 500.0}
|
||||
await _finalize(tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_merges_entries_written_by_another_process(tmp_path, monkeypatch):
|
||||
stats, tasks, settings_dir = _prepare(tmp_path, monkeypatch)
|
||||
stats_file = settings_dir / "stats" / RecipeOpenStats.STATS_FILENAME
|
||||
|
||||
stats.record_open("r1")
|
||||
await _wait_for_save(stats_file)
|
||||
first_ts = json.loads(stats_file.read_text(encoding="utf-8"))["r1"]
|
||||
|
||||
# Another process writes its own entry plus a newer timestamp for r1
|
||||
stats_file.write_text(
|
||||
json.dumps({"r1": first_ts + 100000.0, "r2": 500.0}), encoding="utf-8"
|
||||
)
|
||||
|
||||
stats.record_open("r3")
|
||||
await stats.save_stats(force=True)
|
||||
|
||||
data = json.loads(stats_file.read_text(encoding="utf-8"))
|
||||
# r2 from the other process survives; r1 keeps the newer disk timestamp;
|
||||
# r3 from this process is added
|
||||
assert data["r1"] == first_ts + 100000.0
|
||||
assert data["r2"] == 500.0
|
||||
assert isinstance(data["r3"], float)
|
||||
await _finalize(tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_opened_map_returns_copy(tmp_path, monkeypatch):
|
||||
stats, tasks, _ = _prepare(tmp_path, monkeypatch)
|
||||
stats.record_open("r1")
|
||||
|
||||
opened_map = stats.get_opened_map()
|
||||
opened_map["injected"] = 1.0
|
||||
assert "injected" not in stats.get_opened_map()
|
||||
await _finalize(tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_stats_file_returns_empty_map(tmp_path, monkeypatch):
|
||||
stats, tasks, _ = _prepare(tmp_path, monkeypatch)
|
||||
assert stats.get_opened_map() == {}
|
||||
await _finalize(tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_stats_skips_when_not_dirty(tmp_path, monkeypatch):
|
||||
stats, tasks, settings_dir = _prepare(tmp_path, monkeypatch)
|
||||
stats_file = settings_dir / "stats" / RecipeOpenStats.STATS_FILENAME
|
||||
|
||||
assert await stats.save_stats() is False
|
||||
assert not stats_file.exists()
|
||||
await _finalize(tasks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_ignores_corrupt_file(tmp_path, monkeypatch):
|
||||
settings_dir = tmp_path / "settings"
|
||||
settings_dir.mkdir(parents=True, exist_ok=True)
|
||||
stats_file = settings_dir / "stats" / RecipeOpenStats.STATS_FILENAME
|
||||
stats_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
stats_file.write_text("{not valid json", encoding="utf-8")
|
||||
|
||||
monkeypatch.setattr(
|
||||
stats_module, "get_settings_dir", lambda create=True: str(settings_dir)
|
||||
)
|
||||
RecipeOpenStats._instance = None
|
||||
stats = RecipeOpenStats()
|
||||
assert stats.get_opened_map() == {}
|
||||
Reference in New Issue
Block a user