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