From 1d9c115225779895894cdd579b21c75eb6110885 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Tue, 16 May 2023 10:48:22 +1000
Subject: [PATCH] feat(nodes): add low and high to RandomIntInvocation

---
 invokeai/app/invocations/math.py | 43 +++++++++++++++++++++-----------
 1 file changed, 29 insertions(+), 14 deletions(-)

diff --git a/invokeai/app/invocations/math.py b/invokeai/app/invocations/math.py
index 98f87d2dd4..2ce58c016b 100644
--- a/invokeai/app/invocations/math.py
+++ b/invokeai/app/invocations/math.py
@@ -5,7 +5,12 @@ from typing import Literal
 from pydantic import BaseModel, Field
 import numpy as np
 
-from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext, InvocationConfig
+from .baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    InvocationContext,
+    InvocationConfig,
+)
 
 
 class MathInvocationConfig(BaseModel):
@@ -22,19 +27,21 @@ class MathInvocationConfig(BaseModel):
 
 class IntOutput(BaseInvocationOutput):
     """An integer output"""
-    #fmt: off
+
+    # fmt: off
     type: Literal["int_output"] = "int_output"
     a: int = Field(default=None, description="The output integer")
-    #fmt: on
+    # fmt: on
 
 
 class AddInvocation(BaseInvocation, MathInvocationConfig):
     """Adds two numbers"""
-    #fmt: off
+
+    # fmt: off
     type: Literal["add"] = "add"
     a: int = Field(default=0, description="The first number")
     b: int = Field(default=0, description="The second number")
-    #fmt: on
+    # fmt: on
 
     def invoke(self, context: InvocationContext) -> IntOutput:
         return IntOutput(a=self.a + self.b)
@@ -42,11 +49,12 @@ class AddInvocation(BaseInvocation, MathInvocationConfig):
 
 class SubtractInvocation(BaseInvocation, MathInvocationConfig):
     """Subtracts two numbers"""
-    #fmt: off
+
+    # fmt: off
     type: Literal["sub"] = "sub"
     a: int = Field(default=0, description="The first number")
     b: int = Field(default=0, description="The second number")
-    #fmt: on
+    # fmt: on
 
     def invoke(self, context: InvocationContext) -> IntOutput:
         return IntOutput(a=self.a - self.b)
@@ -54,11 +62,12 @@ class SubtractInvocation(BaseInvocation, MathInvocationConfig):
 
 class MultiplyInvocation(BaseInvocation, MathInvocationConfig):
     """Multiplies two numbers"""
-    #fmt: off
+
+    # fmt: off
     type: Literal["mul"] = "mul"
     a: int = Field(default=0, description="The first number")
     b: int = Field(default=0, description="The second number")
-    #fmt: on
+    # fmt: on
 
     def invoke(self, context: InvocationContext) -> IntOutput:
         return IntOutput(a=self.a * self.b)
@@ -66,11 +75,12 @@ class MultiplyInvocation(BaseInvocation, MathInvocationConfig):
 
 class DivideInvocation(BaseInvocation, MathInvocationConfig):
     """Divides two numbers"""
-    #fmt: off
+
+    # fmt: off
     type: Literal["div"] = "div"
     a: int = Field(default=0, description="The first number")
     b: int = Field(default=0, description="The second number")
-    #fmt: on
+    # fmt: on
 
     def invoke(self, context: InvocationContext) -> IntOutput:
         return IntOutput(a=int(self.a / self.b))
@@ -78,8 +88,13 @@ class DivideInvocation(BaseInvocation, MathInvocationConfig):
 
 class RandomIntInvocation(BaseInvocation):
     """Outputs a single random integer."""
-    #fmt: off
+
+    # fmt: off
     type: Literal["rand_int"] = "rand_int"
-    #fmt: on
+    low: int = Field(default=0, description="The inclusive low value")
+    high: int = Field(
+        default=np.iinfo(np.int32).max, description="The exclusive high value"
+    )
+    # fmt: on
     def invoke(self, context: InvocationContext) -> IntOutput:
-        return IntOutput(a=np.random.randint(0, np.iinfo(np.int32).max))
+        return IntOutput(a=np.random.randint(self.low, self.high))