Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(nodes): add low and high to RandomIntInvocation

Browse Source
This commit is contained in:
psychedelicious 2023-05-16 10:48:22 +10:00
parent 30af20a056
commit 1d9c115225
1 changed files with 29 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
invokeai/app/invocations/math.py
View File

@ -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))