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 from pydantic import BaseModel, Field
import numpy as np import numpy as np
from .baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext, InvocationConfig from .baseinvocation import (
BaseInvocation,
BaseInvocationOutput,
InvocationContext,
InvocationConfig,
)
class MathInvocationConfig(BaseModel): class MathInvocationConfig(BaseModel):
@ -22,19 +27,21 @@ class MathInvocationConfig(BaseModel):
class IntOutput(BaseInvocationOutput): class IntOutput(BaseInvocationOutput):
"""An integer output""" """An integer output"""
#fmt: off
# fmt: off
type: Literal["int_output"] = "int_output" type: Literal["int_output"] = "int_output"
a: int = Field(default=None, description="The output integer") a: int = Field(default=None, description="The output integer")
#fmt: on # fmt: on
class AddInvocation(BaseInvocation, MathInvocationConfig): class AddInvocation(BaseInvocation, MathInvocationConfig):
"""Adds two numbers""" """Adds two numbers"""
#fmt: off
# fmt: off
type: Literal["add"] = "add" type: Literal["add"] = "add"
a: int = Field(default=0, description="The first number") a: int = Field(default=0, description="The first number")
b: int = Field(default=0, description="The second number") b: int = Field(default=0, description="The second number")
#fmt: on # fmt: on
def invoke(self, context: InvocationContext) -> IntOutput: def invoke(self, context: InvocationContext) -> IntOutput:
return IntOutput(a=self.a + self.b) return IntOutput(a=self.a + self.b)
@ -42,11 +49,12 @@ class AddInvocation(BaseInvocation, MathInvocationConfig):
class SubtractInvocation(BaseInvocation, MathInvocationConfig): class SubtractInvocation(BaseInvocation, MathInvocationConfig):
"""Subtracts two numbers""" """Subtracts two numbers"""
#fmt: off
# fmt: off
type: Literal["sub"] = "sub" type: Literal["sub"] = "sub"
a: int = Field(default=0, description="The first number") a: int = Field(default=0, description="The first number")
b: int = Field(default=0, description="The second number") b: int = Field(default=0, description="The second number")
#fmt: on # fmt: on
def invoke(self, context: InvocationContext) -> IntOutput: def invoke(self, context: InvocationContext) -> IntOutput:
return IntOutput(a=self.a - self.b) return IntOutput(a=self.a - self.b)
@ -54,11 +62,12 @@ class SubtractInvocation(BaseInvocation, MathInvocationConfig):
class MultiplyInvocation(BaseInvocation, MathInvocationConfig): class MultiplyInvocation(BaseInvocation, MathInvocationConfig):
"""Multiplies two numbers""" """Multiplies two numbers"""
#fmt: off
# fmt: off
type: Literal["mul"] = "mul" type: Literal["mul"] = "mul"
a: int = Field(default=0, description="The first number") a: int = Field(default=0, description="The first number")
b: int = Field(default=0, description="The second number") b: int = Field(default=0, description="The second number")
#fmt: on # fmt: on
def invoke(self, context: InvocationContext) -> IntOutput: def invoke(self, context: InvocationContext) -> IntOutput:
return IntOutput(a=self.a * self.b) return IntOutput(a=self.a * self.b)
@ -66,11 +75,12 @@ class MultiplyInvocation(BaseInvocation, MathInvocationConfig):
class DivideInvocation(BaseInvocation, MathInvocationConfig): class DivideInvocation(BaseInvocation, MathInvocationConfig):
"""Divides two numbers""" """Divides two numbers"""
#fmt: off
# fmt: off
type: Literal["div"] = "div" type: Literal["div"] = "div"
a: int = Field(default=0, description="The first number") a: int = Field(default=0, description="The first number")
b: int = Field(default=0, description="The second number") b: int = Field(default=0, description="The second number")
#fmt: on # fmt: on
def invoke(self, context: InvocationContext) -> IntOutput: def invoke(self, context: InvocationContext) -> IntOutput:
return IntOutput(a=int(self.a / self.b)) return IntOutput(a=int(self.a / self.b))
@ -78,8 +88,13 @@ class DivideInvocation(BaseInvocation, MathInvocationConfig):
class RandomIntInvocation(BaseInvocation): class RandomIntInvocation(BaseInvocation):
"""Outputs a single random integer.""" """Outputs a single random integer."""
#fmt: off
# fmt: off
type: Literal["rand_int"] = "rand_int" 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: 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))