feat(loaders): add control_after_generate random model selection to checkpoint/unet loaders

The Checkpoint/Unet Loader (LoraManager) nodes now support ComfyUI's
built-in control_after_generate mechanism on the ckpt_name/unet_name combos,
letting users pick a random model on every queue with the selected model
written back into the widget (visible, and lockable via the 'fixed' mode).

A base_model input narrows the random pool: a front-end extension fetches
the name/base_model mapping from the new /api/lm/checkpoints/loader-pool
endpoint and filters the combo options, wired through the node callback,
the refreshComboInNodes extension hook, and a graph.onConfigure hook
installed from onAdded (onNodeCreated fires before the node is attached to
a graph, so the graph reference is unavailable there).

py/nodes/checkpoint_loader.py

@@ -13,6 +13,10 @@ class CheckpointLoaderLM:
 
     Loads checkpoints from both standard ComfyUI folders and LoRA Manager's
     extra folder paths, providing a unified interface for checkpoint loading.
+    The ckpt_name combo supports ComfyUI's control_after_generate, letting
+    users pick a random checkpoint on every run; the base_model input narrows
+    the random pool through a front-end extension that filters the combo
+    options.
     """
 
     NAME = "Checkpoint Loader (LoraManager)"
@@ -22,11 +26,29 @@ class CheckpointLoaderLM:
     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."},
+                    {
+                        "tooltip": (
+                            "The name of the checkpoint (model) to load. Use "
+                            "control_after_generate to pick a random model on "
+                            "every run."
+                        ),
+                        "control_after_generate": True,
+                    },
+                ),
+                "base_model": (
+                    base_models,
+                    {
+                        "default": "Any",
+                        "tooltip": (
+                            "Restrict the random selection pool to this base "
+                            "model. 'Any' uses the full pool."
+                        ),
+                    },
                 ),
             }
         }
@@ -93,15 +115,68 @@ class CheckpointLoaderLM:
             logger.error(f"Error getting checkpoint names: {e}")
             return []
 
-    def load_checkpoint(self, ckpt_name: str) -> Tuple[Any, Any, Any]:
+    @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"]
+
+    @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())
+
+    def load_checkpoint(
+        self, ckpt_name: str, base_model: str = "Any"
+    ) -> Tuple[Any, Any, Any]:
         """Load a checkpoint by name, supporting extra folder paths
 
         Args:
             ckpt_name: The name of the checkpoint to load (relative path with extension)
+            base_model: Only used by the front-end to filter the random pool
 
         Returns:
             Tuple of (MODEL, CLIP, VAE)
         """
+        del base_model
         # Get absolute path from cache using ComfyUI-style name
         ckpt_path, metadata = get_checkpoint_info_absolute(ckpt_name)
 

py/nodes/unet_loader.py

@@ -28,6 +28,10 @@ class UNETLoaderLM:
     Loads diffusion models/UNets from both standard ComfyUI folders and LoRA Manager's
     extra folder paths, providing a unified interface for UNET loading.
     Supports both regular diffusion models and GGUF format models.
+    The unet_name combo supports ComfyUI's control_after_generate, letting
+    users pick a random diffusion model on every run; the base_model input
+    narrows the random pool through a front-end extension that filters the
+    combo options.
     """
 
     NAME = "Unet Loader (LoraManager)"
@@ -37,16 +41,34 @@ class UNETLoaderLM:
     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."},
+                    {
+                        "tooltip": (
+                            "The name of the diffusion model to load. Use "
+                            "control_after_generate to pick a random model on "
+                            "every run."
+                        ),
+                        "control_after_generate": True,
+                    },
                 ),
                 "weight_dtype": (
                     ["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],
                     {"tooltip": "The dtype to use for the model weights."},
                 ),
+                "base_model": (
+                    base_models,
+                    {
+                        "default": "Any",
+                        "tooltip": (
+                            "Restrict the random selection pool to this base "
+                            "model. 'Any' uses the full pool."
+                        ),
+                    },
+                ),
             }
         }
 
@@ -108,16 +130,69 @@ class UNETLoaderLM:
             logger.error(f"Error getting unet names: {e}")
             return []
 
-    def load_unet(self, unet_name: str, weight_dtype: str) -> Tuple[Any, ...]:
+    @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"]
+
+    @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())
+
+    def load_unet(
+        self, unet_name: str, weight_dtype: str, 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
+            base_model: Only used by the front-end to filter the random pool
 
         Returns:
             Tuple of (MODEL,)
         """
+        del base_model
         import torch
 
         # Get absolute path from cache using ComfyUI-style name

py/routes/checkpoint_routes.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from typing import Any, Dict, List, Set
 from aiohttp import web
 
@@ -7,6 +8,7 @@ from .model_route_registrar import ModelRouteRegistrar
 from ..services.checkpoint_service import CheckpointService
 from ..services.service_registry import ServiceRegistry
 from ..config import config
+from ..utils.utils import _format_model_name_for_comfyui
 
 logger = logging.getLogger(__name__)
 
@@ -44,7 +46,45 @@ class CheckpointRoutes(BaseModelRoutes):
         # Checkpoint roots and Unet roots
         registrar.add_prefixed_route('GET', '/api/lm/{prefix}/checkpoints_roots', prefix, self.get_checkpoints_roots)
         registrar.add_prefixed_route('GET', '/api/lm/{prefix}/unet_roots', prefix, self.get_unet_roots)
+
+        # Name/base_model pool for the Random Checkpoint/Unet Loader nodes
+        registrar.add_prefixed_route('GET', '/api/lm/{prefix}/loader-pool', prefix, self.get_loader_pool)
     
+    async def get_loader_pool(self, request: web.Request) -> web.Response:
+        """Return ComfyUI-formatted model names with their base_model.
+
+        Backing data for the Random Checkpoint/Unet Loader nodes: the front-end
+        filters the ckpt_name/unet_name combo options by base_model using this
+        pool, so control_after_generate randomizes within the narrowed set.
+        """
+        try:
+            sub_type = request.query.get("sub_type", "checkpoint")
+            if sub_type not in ("checkpoint", "diffusion_model"):
+                return web.json_response({"error": "invalid sub_type"}, status=400)
+            scanner = await ServiceRegistry.get_checkpoint_scanner()
+            cache = await scanner.get_cached_data()
+            model_roots = scanner.get_model_roots()
+            items: List[Dict[str, str]] = []
+            for item in cache.raw_data:
+                if item.get("sub_type") != sub_type:
+                    continue
+                file_path = item.get("file_path", "")
+                if not file_path or not os.path.exists(file_path):
+                    continue
+                formatted_name = _format_model_name_for_comfyui(file_path, model_roots)
+                if formatted_name:
+                    items.append(
+                        {
+                            "name": formatted_name,
+                            "base_model": item.get("base_model", "") or "",
+                        }
+                    )
+            items.sort(key=lambda x: x["name"])
+            return web.json_response({"items": items})
+        except Exception as e:
+            logger.error(f"Error getting loader pool: {e}", exc_info=True)
+            return web.json_response({"error": str(e)}, status=500)
+
     def _validate_civitai_model_type(self, model_type: str) -> bool:
         """Validate CivitAI model type for Checkpoint"""
         return model_type.lower() == 'checkpoint'