feat(nodes): add RangeOfSizeInvocation

The `RangeInvocation` is a simple wrapper around `range()`, but you must provide `stop > start`.

`RangeOfSizeInvocation` replaces the `stop` parameter with `size`, so that you can just provide the `start` and `step` and get a range of `size` length.
This commit is contained in:
psychedelicious 2023-05-24 15:04:04 +10:00 committed by Kent Keirsey
parent 8f393b64b8
commit b25c1af018

View File

@ -1,6 +1,6 @@
# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
# Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654) and the InvokeAI Team
from typing import Literal, Optional
from typing import Literal
import numpy as np
from pydantic import Field, validator
@ -24,7 +24,7 @@ class IntCollectionOutput(BaseInvocationOutput):
class RangeInvocation(BaseInvocation):
"""Creates a range"""
"""Creates a range of numbers from start to stop with step"""
type: Literal["range"] = "range"
@ -45,6 +45,22 @@ class RangeInvocation(BaseInvocation):
)
class RangeOfSizeInvocation(BaseInvocation):
"""Creates a range from start to start + size with step"""
type: Literal["range_of_size"] = "range_of_size"
# Inputs
start: int = Field(default=0, description="The start of the range")
size: int = Field(default=1, description="The number of values")
step: int = Field(default=1, description="The step of the range")
def invoke(self, context: InvocationContext) -> IntCollectionOutput:
return IntCollectionOutput(
collection=list(range(self.start, self.start + self.size + 1, self.step))
)
class RandomRangeInvocation(BaseInvocation):
"""Creates a collection of random numbers"""