Add an activity field to node progress and report model loading

A node that is loading weights looks identical to one that is computing:
progress_state carries only state/value/max, and load_models_gpu runs inside
the sampler, so the UI shows the sampler sitting at 0% while weights move.

Adds an optional activity field to the node progress state, reported through a
hook like the existing progress bar hook, and sets it to "loading" around the
VRAM load. The field is absent unless something sets it, so existing clients
are unaffected.

comfy/model_management.py

@@ -1017,7 +1017,8 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
         if vram_set_state == VRAMState.NO_VRAM:
             lowvram_model_memory = 0.1
 
-        loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)
+        with comfy.utils.progress_activity("loading"):
+            loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)
         vram_used = 0 if is_device_cpu(torch_dev) else loaded_model.model_loaded_memory()
         ram_used = model.loaded_ram_size() if model.is_dynamic() else loaded_model.model_memory() - vram_used
         detail("Model loaded: patcher=%s model=%s ram_mb=%.1f vram_mb=%.1f", model.__class__.__name__, model.model.__class__.__name__, ram_used / (1024 ** 2), vram_used / (1024 ** 2))

comfy/utils.py

@@ -22,6 +22,7 @@ import math
 import struct
 import ctypes
 import os
+import contextlib
 import comfy.memory_management
 import safetensors.torch
 import numpy as np
@@ -1297,6 +1298,26 @@ def set_progress_bar_global_hook(function):
     global PROGRESS_BAR_HOOK
     PROGRESS_BAR_HOOK = function
 
+PROGRESS_ACTIVITY_HOOK = None
+def set_progress_activity_global_hook(function):
+    global PROGRESS_ACTIVITY_HOOK
+    PROGRESS_ACTIVITY_HOOK = function
+
+@contextlib.contextmanager
+def progress_activity(activity):
+    """Report what the running node is spending time on, so the UI can tell a
+    slow load apart from slow compute."""
+    hook = PROGRESS_ACTIVITY_HOOK
+    if hook is None:
+        yield
+        return
+
+    hook(activity)
+    try:
+        yield
+    finally:
+        hook(None)
+
 # Throttle settings for progress bar updates to reduce WebSocket flooding
 PROGRESS_THROTTLE_MIN_INTERVAL = 0.1  # 100ms minimum between updates
 PROGRESS_THROTTLE_MIN_PERCENT = 0.5   # 0.5% minimum progress change

comfy_execution/progress.py

@@ -1,5 +1,5 @@
 from typing import TypedDict, Dict, Optional, Tuple
-from typing_extensions import override
+from typing_extensions import override, NotRequired
 from PIL import Image
 from enum import Enum
 from abc import ABC
@@ -27,6 +27,7 @@ class NodeProgressState(TypedDict):
     state: NodeState
     value: float
     max: float
+    activity: NotRequired[str]  # what a running node is spending time on, e.g. "loading"
 
 
 class ProgressHandler(ABC):
@@ -163,8 +164,11 @@ class WebUIProgressHandler(ProgressHandler):
             return
 
         # Only send info for non-pending nodes
-        active_nodes = {
-            node_id: {
+        active_nodes = {}
+        for node_id, state in nodes.items():
+            if state["state"] == NodeState.Pending:
+                continue
+            active_nodes[node_id] = {
                 "value": state["value"],
                 "max": state["max"],
                 "state": state["state"].value,
@@ -174,9 +178,8 @@ class WebUIProgressHandler(ProgressHandler):
                 "parent_node_id": self.registry.dynprompt.get_parent_node_id(node_id),
                 "real_node_id": self.registry.dynprompt.get_real_node_id(node_id),
             }
-            for node_id, state in nodes.items()
-            if state["state"] != NodeState.Pending
-        }
+            if "activity" in state:
+                active_nodes[node_id]["activity"] = state["activity"]
 
         # Send a combined progress_state message with all node states
         # Include client_id to ensure message is only sent to the initiating client
@@ -301,6 +304,23 @@ class ProgressRegistry:
                     node_id, value, max_value, entry, self.prompt_id, image
                 )
 
+    def set_activity(self, node_id: str, activity: str | None) -> None:
+        """Set what a running node is spending time on, or None to clear it"""
+        entry = self.ensure_entry(node_id)
+        if entry.get("activity") == activity:
+            return
+        if activity is None:
+            entry.pop("activity", None)
+        else:
+            entry["activity"] = activity
+
+        # Notify all enabled handlers
+        for handler in self.handlers.values():
+            if handler.enabled:
+                handler.update_handler(
+                    node_id, entry["value"], entry["max"], entry, self.prompt_id
+                )
+
     def finish_progress(self, node_id: str) -> None:
         """Finish progress tracking for a node"""
         entry = self.ensure_entry(node_id)

main.py

@@ -483,7 +483,14 @@ def hijack_progress(server_instance):
                     server_instance.client_id,
                 )
 
+    def activity_hook(activity):
+        executing_context = get_executing_context()
+        if executing_context is None:
+            return
+        get_progress_state().set_activity(executing_context.node_id, activity)
+
     comfy.utils.set_progress_bar_global_hook(hook)
+    comfy.utils.set_progress_activity_global_hook(activity_hook)
 
 
 def cleanup_temp():

tests-unit/execution_test/progress_activity_test.py

@@ -0,0 +1,90 @@
+from unittest.mock import Mock
+
+import comfy.utils
+from comfy_execution.progress import (
+    NodeState,
+    ProgressRegistry,
+    WebUIProgressHandler,
+)
+
+
+def make_registry():
+    dynprompt = Mock()
+    dynprompt.get_display_node_id.return_value = "1"
+    dynprompt.get_parent_node_id.return_value = None
+    dynprompt.get_real_node_id.return_value = "1"
+    return ProgressRegistry(prompt_id="p", dynprompt=dynprompt)
+
+
+def test_activity_is_set_and_cleared():
+    registry = make_registry()
+    registry.start_progress("1")
+
+    registry.set_activity("1", "loading")
+    assert registry.nodes["1"]["activity"] == "loading"
+
+    registry.set_activity("1", None)
+    assert "activity" not in registry.nodes["1"]
+
+
+def test_unchanged_activity_does_not_notify():
+    registry = make_registry()
+    handler = Mock(enabled=True)
+    registry.handlers["test"] = handler
+    registry.start_progress("1")
+
+    registry.set_activity("1", "loading")
+    assert handler.update_handler.call_count == 1
+
+    registry.set_activity("1", "loading")
+    assert handler.update_handler.call_count == 1
+
+    registry.set_activity("1", None)
+    assert handler.update_handler.call_count == 2
+
+
+def test_activity_is_only_sent_when_set():
+    server = Mock(client_id="c")
+    registry = make_registry()
+    handler = WebUIProgressHandler(server)
+    handler.set_registry(registry)
+    registry.register_handler(handler)
+
+    registry.start_progress("1")
+    nodes = server.send_sync.call_args[0][1]["nodes"]
+    assert "activity" not in nodes["1"]
+
+    registry.set_activity("1", "loading")
+    nodes = server.send_sync.call_args[0][1]["nodes"]
+    assert nodes["1"]["activity"] == "loading"
+    assert nodes["1"]["state"] == NodeState.Running.value
+
+
+def test_progress_activity_reports_and_clears():
+    seen = []
+    comfy.utils.set_progress_activity_global_hook(seen.append)
+    try:
+        with comfy.utils.progress_activity("loading"):
+            assert seen == ["loading"]
+    finally:
+        comfy.utils.set_progress_activity_global_hook(None)
+    assert seen == ["loading", None]
+
+
+def test_progress_activity_clears_on_error():
+    seen = []
+    comfy.utils.set_progress_activity_global_hook(seen.append)
+    try:
+        with comfy.utils.progress_activity("loading"):
+            raise RuntimeError("boom")
+    except RuntimeError:
+        pass
+    finally:
+        comfy.utils.set_progress_activity_global_hook(None)
+    assert seen == ["loading", None]
+
+
+def test_progress_activity_without_hook():
+    comfy.utils.set_progress_activity_global_hook(None)
+    with comfy.utils.progress_activity("loading"):
+        pass