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
Download Patch File Download Diff File Expand all files Collapse all files

19
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,6 +27,7 @@ class MathInvocationConfig(BaseModel):
class IntOutput(BaseInvocationOutput):
"""An integer output"""
# fmt: off
type: Literal["int_output"] = "int_output"
a: int = Field(default=None, description="The output integer")
@ -30,6 +36,7 @@ class IntOutput(BaseInvocationOutput):
class AddInvocation(BaseInvocation, MathInvocationConfig):
"""Adds two numbers"""
# fmt: off
type: Literal["add"] = "add"
a: int = Field(default=0, description="The first number")
@ -42,6 +49,7 @@ class AddInvocation(BaseInvocation, MathInvocationConfig):
class SubtractInvocation(BaseInvocation, MathInvocationConfig):
"""Subtracts two numbers"""
# fmt: off
type: Literal["sub"] = "sub"
a: int = Field(default=0, description="The first number")
@ -54,6 +62,7 @@ class SubtractInvocation(BaseInvocation, MathInvocationConfig):
class MultiplyInvocation(BaseInvocation, MathInvocationConfig):
"""Multiplies two numbers"""
# fmt: off
type: Literal["mul"] = "mul"
a: int = Field(default=0, description="The first number")
@ -66,6 +75,7 @@ class MultiplyInvocation(BaseInvocation, MathInvocationConfig):
class DivideInvocation(BaseInvocation, MathInvocationConfig):
"""Divides two numbers"""
# fmt: off
type: Literal["div"] = "div"
a: int = Field(default=0, description="The first number")
@ -78,8 +88,13 @@ class DivideInvocation(BaseInvocation, MathInvocationConfig):
class RandomIntInvocation(BaseInvocation):
"""Outputs a single random integer."""
# fmt: off
type: Literal["rand_int"] = "rand_int"
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))