From 763debdeebeb80133cb1be79fb57f279d5513292 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Mon, 19 Feb 2024 12:57:05 +1100
Subject: [PATCH] fix(nodes): fix typing on stats service context manager

---
 .../app/services/invocation_stats/invocation_stats_base.py    | 4 ++--
 .../app/services/invocation_stats/invocation_stats_default.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/invokeai/app/services/invocation_stats/invocation_stats_base.py b/invokeai/app/services/invocation_stats/invocation_stats_base.py
index b28220e74c..3266d985fe 100644
--- a/invokeai/app/services/invocation_stats/invocation_stats_base.py
+++ b/invokeai/app/services/invocation_stats/invocation_stats_base.py
@@ -30,7 +30,7 @@ writes to the system log is stored in InvocationServices.performance_statistics.
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Iterator
+from typing import ContextManager
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation
 from invokeai.app.services.invocation_stats.invocation_stats_common import InvocationStatsSummary
@@ -50,7 +50,7 @@ class InvocationStatsServiceBase(ABC):
         self,
         invocation: BaseInvocation,
         graph_execution_state_id: str,
-    ) -> Iterator[None]:
+    ) -> ContextManager[None]:
         """
         Return a context object that will capture the statistics on the execution
         of invocaation. Use with: to place around the part of the code that executes the invocation.
diff --git a/invokeai/app/services/invocation_stats/invocation_stats_default.py b/invokeai/app/services/invocation_stats/invocation_stats_default.py
index 06a5b675c3..5a41f1f5d6 100644
--- a/invokeai/app/services/invocation_stats/invocation_stats_default.py
+++ b/invokeai/app/services/invocation_stats/invocation_stats_default.py
@@ -2,7 +2,7 @@ import json
 import time
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Iterator
+from typing import Generator
 
 import psutil
 import torch
@@ -41,7 +41,7 @@ class InvocationStatsService(InvocationStatsServiceBase):
         self._invoker = invoker
 
     @contextmanager
-    def collect_stats(self, invocation: BaseInvocation, graph_execution_state_id: str) -> Iterator[None]:
+    def collect_stats(self, invocation: BaseInvocation, graph_execution_state_id: str) -> Generator[None, None, None]:
         # This is to handle case of the model manager not being initialized, which happens
         # during some tests.
         services = self._invoker.services