From 0fdcc0af658a5dbbed292f8e5658f4c93dfc9afc Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Fri, 1 Dec 2023 07:41:01 +1100
Subject: [PATCH] feat(nodes): add index and total to iterate output

---
 invokeai/app/services/shared/graph.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py
index c825a84011..854defc945 100644
--- a/invokeai/app/services/shared/graph.py
+++ b/invokeai/app/services/shared/graph.py
@@ -207,10 +207,12 @@ class IterateInvocationOutput(BaseInvocationOutput):
     item: Any = OutputField(
         description="The item being iterated over", title="Collection Item", ui_type=UIType._CollectionItem
     )
+    index: int = OutputField(description="The index of the item", title="Index")
+    total: int = OutputField(description="The total number of items", title="Total")
 
 
 # TODO: Fill this out and move to invocations
-@invocation("iterate", version="1.0.0")
+@invocation("iterate", version="1.1.0")
 class IterateInvocation(BaseInvocation):
     """Iterates over a list of items"""
 
@@ -221,7 +223,7 @@ class IterateInvocation(BaseInvocation):
 
     def invoke(self, context: InvocationContext) -> IterateInvocationOutput:
         """Produces the outputs as values"""
-        return IterateInvocationOutput(item=self.collection[self.index])
+        return IterateInvocationOutput(item=self.collection[self.index], index=self.index, total=len(self.collection))
 
 
 @invocation_output("collect_output")