feat(nodes): WIP restricted invocation context #4910
Conversation
Upgrade pydantic and fastapi to latest.
- pydantic~=2.4.2
- fastapi~=103.2
- fastapi-events~=0.9.1
**Big Changes**
There are a number of logic changes needed to support pydantic v2. Most changes are very simple, like using the new methods to serialized and deserialize models, but there are a few more complex changes.
**Invocations**
The biggest change relates to invocation creation, instantiation and validation.
Because pydantic v2 moves all validation logic into the rust pydantic-core, we may no longer directly stick our fingers into the validation pie.
Previously, we (ab)used models and fields to allow invocation fields to be optional at instantiation, but required when `invoke()` is called. We directly manipulated the fields and invocation models when calling `invoke()`.
With pydantic v2, this is much more involved. Changes to the python wrapper do not propagate down to the rust validation logic - you have to rebuild the model. This causes problem with concurrent access to the invocation classes and is not a free operation.
This logic has been totally refactored and we do not need to change the model any more. The details are in `baseinvocation.py`, in the `InputField` function and `BaseInvocation.invoke_internal()` method.
In the end, this implementation is cleaner.
**Invocation Fields**
In pydantic v2, you can no longer directly add or remove fields from a model.
Previously, we did this to add the `type` field to invocations.
**Invocation Decorators**
With pydantic v2, we instead use the imperative `create_model()` API to create a new model with the additional field. This is done in `baseinvocation.py` in the `invocation()` wrapper.
A similar technique is used for `invocation_output()`.
**Minor Changes**
There are a number of minor changes around the pydantic v2 models API.
**Protected `model_` Namespace**
All models' pydantic-provided methods and attributes are prefixed with `model_` and this is considered a protected namespace. This causes some conflict, because "model" means something to us, and we have a ton of pydantic models with attributes starting with "model_".
Forunately, there are no direct conflicts. However, in any pydantic model where we define an attribute or method that starts with "model_", we must tell set the protected namespaces to an empty tuple.
```py
class IPAdapterModelField(BaseModel):
model_name: str = Field(description="Name of the IP-Adapter model")
base_model: BaseModelType = Field(description="Base model")
model_config = ConfigDict(protected_namespaces=())
```
**Model Serialization**
Pydantic models no longer have `Model.dict()` or `Model.json()`.
Instead, we use `Model.model_dump()` or `Model.model_dump_json()`.
**Model Deserialization**
Pydantic models no longer have `Model.parse_obj()` or `Model.parse_raw()`, and there are no `parse_raw_as()` or `parse_obj_as()` functions.
Instead, you need to create a `TypeAdapter` object to parse python objects or JSON into a model.
```py
adapter_graph = TypeAdapter(Graph)
deserialized_graph_from_json = adapter_graph.validate_json(graph_json)
deserialized_graph_from_dict = adapter_graph.validate_python(graph_dict)
```
**Field Customisation**
Pydantic `Field`s no longer accept arbitrary args.
Now, you must put all additional arbitrary args in a `json_schema_extra` arg on the field.
**Schema Customisation**
FastAPI and pydantic schema generation now follows the OpenAPI version 3.1 spec.
This necessitates two changes:
- Our schema customization logic has been revised
- Schema parsing to build node templates has been revised
The specific aren't important, but this does present additional surface area for bugs.
**Performance Improvements**
Pydantic v2 is a full rewrite with a rust backend. This offers a substantial performance improvement (pydantic claims 5x to 50x depending on the task). We'll notice this the most during serialization and deserialization of sessions/graphs, which happens very very often - a couple times per node.
I haven't done any benchmarks, but anecdotally, graph execution is much faster. Also, very larges graphs - like with massive iterators - are much, much faster.
psychedelicious
force-pushed
the
chore/pydantic-fastapi-upgrade
branch
from
0aedd6d to
d9a2249
Compare
|
Has this been updated with pydantic changes? |
|
Overall, I love this initiative! Thanks for putting together a draft. It's a little hard to review right now because the diff is polluted with pydantic changes. I'll just leave a few high-level comments for now. |
In my opinion, we should aim to have all nodes (core and community) interface with the same API. It's nice to have a clear boundary between nodes and the app backend. I think there are a few downsides to allowing special cases for core nodes:
Put another way, before adding a method to the context we should ask ourselves: "Are we willing to support this method for both core and community nodes indefinitely?" Of course, my mind could be changed if there was a compelling use case. |
What if we defined a standard import path for all node utilities? (e.g. I don't love the idea that the context would be responsible for managing types as well. |
|
I think there would be value in having a single file that node developers can reference that cleanly contains the API that they are expected to interact with. (TBD whether this should just be an ABC or the actual implementation). If we went with an ABC, then in my mind it would look something like this: class InvocationContext(ABC):
@abstractmethod
def get_image(self, name: str) -> ImageType:
"""<insert a very detailed docstring>
"""
pass
@abstractmethod
def save_image(self, name: str) -> ImageType:
"""<insert a very detailed docstring>
"""
pass
...Separately, I think there is an opportunity to simplify the If we were open to re-designing the node API, I think there are other conversations to be had around what this API should look like ( |
|
As somebody interested in color spaces and wider gamuts and how those could be used in the future, I'd like to see |
|
Superseded by #5491 |
What type of PR is this? (check all applicable)
Description
Note: this PR is based on the #4758 branch
feat(nodes): WIP restricted invocation context
This is a first draft of a restricted
InvocationContext, which is provided to all nodes.Motivation
The existing
InvocationContextclass has two issues.1. Too powerful
It's entirely possible that a node could wreck an installation by accidentally calling a destructive method or something like that.
Also, the application relies on nodes to be well-behaved and never do certain things (like delete images), so we should just enforce this explicitly.
It's probably not possible to prevent a node from doing something naughty if it wants to, but we can reduce the chance of unintentional footgunning to almost zero by providing a restricted API.
2. Too low-level/complicated
Many of the service methods are low-level and present an awkward API.
Consider the method to save an image:
There's a crapload of fluff here that nodes shouldn't need to worry about. It's mostly boilerplate.
Draft API
The full-fat context is now called
AppInvocationContext, and the restricted one isInvocationContext. These names can be changed, I just did this so I didn't have to change a bazillion imports in all nodes while testing.See
base_invocation.pyfor the full implementation - it's a simple wrapper.Example improved method:
The above
create()method is reduced to this:The other methods don't have as drastic an improvement but they are all improved.
Discussion points
In addition to the design and implementation, there are a few points that probably warrant discussion.
This restricted API is currently set up to apply to all nodes - core, InvokeAI-managed nodes and future community nodes.
Do we want to maybe let core nodes get access to the original, full-fat context, and only give the restricted context to community nodes?
So far, I've implemented most of what our core nodes use. What if community nodes want more access - say, to star images? How do we decide what makes the cut?
There are some constants/enums/things that we want to provide to nodes. We can do that via the context, as opposed to via imports.
Example:
Images have categories, which are selected via the
ImageCategoryenum - see thecreate()method above for how it works in the current release.In this PR, I've exposed that enum as
context.categoriesand made this an optional arg:This feels neat and tidy but I'm unsure if it's a good idea. Maybe the context should be methods and context attributes only?
Related Tickets & Documents
QA Instructions, Screenshots, Recordings
You can test this now, using a simple txt2img workflow with sd1.5 (other nodes are not migrated). You'll first need to run
pip install -e ".[dev,test]"to get pydantic v2 installed.simple t2i.json
Added/updated tests?
have not been included
[optional] Are there any post deployment tasks we need to perform?