feat(nodes): update float_math and integer_math to use new ui_choice_labels

This commit is contained in:
psychedelicious 2023-09-13 16:24:46 +10:00
parent 57ebf735e6
commit fb188ce63e

View File

@ -103,18 +103,31 @@ class RoundInvocation(BaseInvocation):
INTEGER_OPERATIONS = Literal[ INTEGER_OPERATIONS = Literal[
"Add A+B", "ADD",
"Subtract A-B", "SUB",
"Multiply A*B", "MUL",
"Divide A/B", "DIV",
"Exponentiate A^B", "EXP",
"Modulus A%B", "MOD",
"Absolute Value of A", "ABS",
"Minimum(A,B)", "MIN",
"Maximum(A,B)", "MAX",
] ]
INTEGER_OPERATIONS_LABELS = dict(
ADD="Add A+B",
SUB="Subtract A-B",
MUL="Multiply A*B",
DIV="Divide A/B",
EXP="Exponentiate A^B",
MOD="Modulus A%B",
ABS="Absolute Value of A",
MIN="Minimum(A,B)",
MAX="Maximum(A,B)",
)
@invocation( @invocation(
"integer_math", "integer_math",
title="Integer Math", title="Integer Math",
@ -137,54 +150,70 @@ INTEGER_OPERATIONS = Literal[
class IntegerMathInvocation(BaseInvocation): class IntegerMathInvocation(BaseInvocation):
"""Performs integer math.""" """Performs integer math."""
operation: INTEGER_OPERATIONS = InputField(default="Add A+B", description="The operation to perform") operation: INTEGER_OPERATIONS = InputField(
default="ADD", description="The operation to perform", ui_choice_labels=INTEGER_OPERATIONS_LABELS
)
a: int = InputField(default=0, description=FieldDescriptions.num_1) a: int = InputField(default=0, description=FieldDescriptions.num_1)
b: int = InputField(default=0, description=FieldDescriptions.num_2) b: int = InputField(default=0, description=FieldDescriptions.num_2)
@validator("b") @validator("b")
def no_unrepresentable_results(cls, v, values): def no_unrepresentable_results(cls, v, values):
if values["operation"] == "Divide A/B" and v == 0: if values["operation"] == "DIV" and v == 0:
raise ValueError("Cannot divide by zero") raise ValueError("Cannot divide by zero")
elif values["operation"] == "Modulus A%B" and v == 0: elif values["operation"] == "MOD" and v == 0:
raise ValueError("Cannot divide by zero") raise ValueError("Cannot divide by zero")
elif values["operation"] == "Exponentiate A^B" and v < 0: elif values["operation"] == "EXP" and v < 0:
raise ValueError("Result of exponentiation is not an integer") raise ValueError("Result of exponentiation is not an integer")
return v return v
def invoke(self, context: InvocationContext) -> IntegerOutput: def invoke(self, context: InvocationContext) -> IntegerOutput:
# Python doesn't support switch statements until 3.10, but InvokeAI supports back to 3.9 # Python doesn't support switch statements until 3.10, but InvokeAI supports back to 3.9
if self.operation == "Add A+B": if self.operation == "ADD":
return IntegerOutput(value=self.a + self.b) return IntegerOutput(value=self.a + self.b)
elif self.operation == "Subtract A-B": elif self.operation == "SUB":
return IntegerOutput(value=self.a - self.b) return IntegerOutput(value=self.a - self.b)
elif self.operation == "Multiply A*B": elif self.operation == "MUL":
return IntegerOutput(value=self.a * self.b) return IntegerOutput(value=self.a * self.b)
elif self.operation == "Divide A/B": elif self.operation == "DIV":
return IntegerOutput(value=int(self.a / self.b)) return IntegerOutput(value=int(self.a / self.b))
elif self.operation == "Exponentiate A^B": elif self.operation == "EXP":
return IntegerOutput(value=self.a**self.b) return IntegerOutput(value=self.a**self.b)
elif self.operation == "Modulus A%B": elif self.operation == "MOD":
return IntegerOutput(value=self.a % self.b) return IntegerOutput(value=self.a % self.b)
elif self.operation == "Absolute Value of A": elif self.operation == "ABS":
return IntegerOutput(value=abs(self.a)) return IntegerOutput(value=abs(self.a))
elif self.operation == "Minimum(A,B)": elif self.operation == "MIN":
return IntegerOutput(value=min(self.a, self.b)) return IntegerOutput(value=min(self.a, self.b))
else: # self.operation == "Maximum(A,B)": else: # self.operation == "MAX":
return IntegerOutput(value=max(self.a, self.b)) return IntegerOutput(value=max(self.a, self.b))
FLOAT_OPERATIONS = Literal[ FLOAT_OPERATIONS = Literal[
"Add A+B", "ADD",
"Subtract A-B", "SUB",
"Multiply A*B", "MUL",
"Divide A/B", "DIV",
"Exponentiate A^B", "EXP",
"Absolute Value of A", "ABS",
"Minimum(A,B)", "SQRT",
"Maximum(A,B)", "MIN",
"MAX",
] ]
FLOAT_OPERATIONS_LABELS = dict(
ADD="Add A+B",
SUB="Subtract A-B",
MUL="Multiply A*B",
DIV="Divide A/B",
EXP="Exponentiate A^B",
ABS="Absolute Value of A",
SQRT="Square Root of A",
MIN="Minimum(A,B)",
MAX="Maximum(A,B)",
)
@invocation( @invocation(
"float_math", "float_math",
title="Float Math", title="Float Math",
@ -195,37 +224,39 @@ FLOAT_OPERATIONS = Literal[
class FloatMathInvocation(BaseInvocation): class FloatMathInvocation(BaseInvocation):
"""Performs floating point math.""" """Performs floating point math."""
operation: FLOAT_OPERATIONS = InputField(default="Add A+B", description="The operation to perform") operation: FLOAT_OPERATIONS = InputField(
default="ADD", description="The operation to perform", ui_choice_labels=FLOAT_OPERATIONS_LABELS
)
a: float = InputField(default=0, description=FieldDescriptions.num_1) a: float = InputField(default=0, description=FieldDescriptions.num_1)
b: float = InputField(default=0, description=FieldDescriptions.num_2) b: float = InputField(default=0, description=FieldDescriptions.num_2)
@validator("b") @validator("b")
def no_unrepresentable_results(cls, v, values): def no_unrepresentable_results(cls, v, values):
if values["operation"] == "Divide A/B" and v == 0: if values["operation"] == "DIV" and v == 0:
raise ValueError("Cannot divide by zero") raise ValueError("Cannot divide by zero")
elif values["operation"] == "Exponentiate A^B" and values["a"] == 0 and v < 0: elif values["operation"] == "EXP" and values["a"] == 0 and v < 0:
raise ValueError("Cannot raise zero to a negative power") raise ValueError("Cannot raise zero to a negative power")
elif values["operation"] == "Exponentiate A^B" and type(values["a"] ** v) == complex: elif values["operation"] == "EXP" and type(values["a"] ** v) == complex:
raise ValueError("Root operation resulted in a complex number") raise ValueError("Root operation resulted in a complex number")
return v return v
def invoke(self, context: InvocationContext) -> FloatOutput: def invoke(self, context: InvocationContext) -> FloatOutput:
# Python doesn't support switch statements until 3.10, but InvokeAI supports back to 3.9 # Python doesn't support switch statements until 3.10, but InvokeAI supports back to 3.9
if self.operation == "Add A+B": if self.operation == "ADD":
return FloatOutput(value=self.a + self.b) return FloatOutput(value=self.a + self.b)
elif self.operation == "Subtract A-B": elif self.operation == "SUB":
return FloatOutput(value=self.a - self.b) return FloatOutput(value=self.a - self.b)
elif self.operation == "Multiply A*B": elif self.operation == "MUL":
return FloatOutput(value=self.a * self.b) return FloatOutput(value=self.a * self.b)
elif self.operation == "Divide A/B": elif self.operation == "DIV":
return FloatOutput(value=self.a / self.b) return FloatOutput(value=self.a / self.b)
elif self.operation == "Exponentiate A^B": elif self.operation == "EXP":
return FloatOutput(value=self.a**self.b) return FloatOutput(value=self.a**self.b)
elif self.operation == "Square Root of A": elif self.operation == "SQRT":
return FloatOutput(value=np.sqrt(self.a)) return FloatOutput(value=np.sqrt(self.a))
elif self.operation == "Absolute Value of A": elif self.operation == "ABS":
return FloatOutput(value=abs(self.a)) return FloatOutput(value=abs(self.a))
elif self.operation == "Minimum(A,B)": elif self.operation == "MIN":
return FloatOutput(value=min(self.a, self.b)) return FloatOutput(value=min(self.a, self.b))
else: # self.operation == "Maximum(A,B)": else: # self.operation == "MAX":
return FloatOutput(value=max(self.a, self.b)) return FloatOutput(value=max(self.a, self.b))