2023-07-15 15:06:50 +00:00
|
|
|
# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654) & the InvokeAI Team
|
2023-08-17 22:45:25 +00:00
|
|
|
from typing import Literal
|
2023-03-03 06:02:00 +00:00
|
|
|
|
2023-11-27 10:22:31 +00:00
|
|
|
import cv2
|
2023-07-15 15:06:50 +00:00
|
|
|
import numpy as np
|
|
|
|
|
from PIL import Image
|
feat(api): chore: pydantic & fastapi upgrade
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.
2023-09-24 08:11:07 +00:00
|
|
|
from pydantic import ConfigDict
|
2023-03-03 06:02:00 +00:00
|
|
|
|
2024-01-13 12:23:16 +00:00
|
|
|
from invokeai.app.invocations.fields import ImageField
|
|
|
|
|
from invokeai.app.invocations.primitives import ImageOutput
|
2024-02-05 06:16:35 +00:00
|
|
|
from invokeai.app.services.shared.invocation_context import InvocationContext
|
2024-02-10 13:11:33 +00:00
|
|
|
from invokeai.backend.image_util.basicsr.rrdbnet_arch import RRDBNet
|
2023-11-27 20:37:39 +00:00
|
|
|
from invokeai.backend.image_util.realesrgan.realesrgan import RealESRGAN
|
2022-12-01 05:33:20 +00:00
|
|
|
|
2024-01-13 12:23:16 +00:00
|
|
|
from .baseinvocation import BaseInvocation, invocation
|
2024-02-07 05:33:55 +00:00
|
|
|
from .fields import InputField, WithBoard, WithMetadata
|
2022-12-01 05:33:20 +00:00
|
|
|
|
2023-07-15 15:06:50 +00:00
|
|
|
# TODO: Populate this from disk?
|
|
|
|
|
# TODO: Use model manager to load?
|
2023-07-17 11:00:22 +00:00
|
|
|
ESRGAN_MODELS = Literal[
|
2023-07-15 15:06:50 +00:00
|
|
|
"RealESRGAN_x4plus.pth",
|
|
|
|
|
"RealESRGAN_x4plus_anime_6B.pth",
|
|
|
|
|
"ESRGAN_SRx4_DF2KOST_official-ff704c30.pth",
|
2023-07-17 11:00:22 +00:00
|
|
|
"RealESRGAN_x2plus.pth",
|
2023-07-15 15:06:50 +00:00
|
|
|
]
|
2023-05-24 05:50:55 +00:00
|
|
|
|
2024-03-19 01:11:13 +00:00
|
|
|
ESRGAN_MODEL_URLS: dict[str, str] = {
|
|
|
|
|
"RealESRGAN_x4plus.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
|
|
|
|
|
"RealESRGAN_x4plus_anime_6B.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
|
2024-03-22 05:17:45 +00:00
|
|
|
"ESRGAN_SRx4_DF2KOST_official-ff704c30.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth",
|
2024-03-19 01:11:13 +00:00
|
|
|
"RealESRGAN_x2plus.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
|
|
|
|
|
}
|
|
|
|
|
|
2022-12-01 05:33:20 +00:00
|
|
|
|
2024-03-19 11:08:16 +00:00
|
|
|
@invocation("esrgan", title="Upscale (RealESRGAN)", tags=["esrgan", "upscale"], category="esrgan", version="1.3.2")
|
2024-02-07 05:33:55 +00:00
|
|
|
class ESRGANInvocation(BaseInvocation, WithMetadata, WithBoard):
|
2023-07-15 15:06:50 +00:00
|
|
|
"""Upscales an image using RealESRGAN."""
|
2023-04-10 09:07:48 +00:00
|
|
|
|
2023-08-14 03:23:09 +00:00
|
|
|
image: ImageField = InputField(description="The input image")
|
|
|
|
|
model_name: ESRGAN_MODELS = InputField(default="RealESRGAN_x4plus.pth", description="The Real-ESRGAN model to use")
|
2023-10-04 01:43:16 +00:00
|
|
|
tile_size: int = InputField(
|
2023-10-04 04:23:31 +00:00
|
|
|
default=400, ge=0, description="Tile size for tiled ESRGAN upscaling (0=tiling disabled)"
|
2023-10-04 01:43:16 +00:00
|
|
|
)
|
2023-07-18 14:26:45 +00:00
|
|
|
|
feat(api): chore: pydantic & fastapi upgrade
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.
2023-09-24 08:11:07 +00:00
|
|
|
model_config = ConfigDict(protected_namespaces=())
|
|
|
|
|
|
2024-02-05 06:16:35 +00:00
|
|
|
def invoke(self, context: InvocationContext) -> ImageOutput:
|
2024-01-13 12:23:16 +00:00
|
|
|
image = context.images.get_pil(self.image.image_name)
|
2023-07-15 15:06:50 +00:00
|
|
|
|
|
|
|
|
rrdbnet_model = None
|
|
|
|
|
netscale = None
|
|
|
|
|
|
|
|
|
|
if self.model_name in [
|
|
|
|
|
"RealESRGAN_x4plus.pth",
|
|
|
|
|
"ESRGAN_SRx4_DF2KOST_official-ff704c30.pth",
|
|
|
|
|
]:
|
|
|
|
|
# x4 RRDBNet model
|
|
|
|
|
rrdbnet_model = RRDBNet(
|
|
|
|
|
num_in_ch=3,
|
|
|
|
|
num_out_ch=3,
|
|
|
|
|
num_feat=64,
|
|
|
|
|
num_block=23,
|
|
|
|
|
num_grow_ch=32,
|
|
|
|
|
scale=4,
|
|
|
|
|
)
|
|
|
|
|
netscale = 4
|
2023-07-16 00:54:52 +00:00
|
|
|
elif self.model_name in ["RealESRGAN_x4plus_anime_6B.pth"]:
|
2023-07-15 15:06:50 +00:00
|
|
|
# x4 RRDBNet model, 6 blocks
|
|
|
|
|
rrdbnet_model = RRDBNet(
|
|
|
|
|
num_in_ch=3,
|
|
|
|
|
num_out_ch=3,
|
|
|
|
|
num_feat=64,
|
2023-07-16 00:54:52 +00:00
|
|
|
num_block=6, # 6 blocks
|
2023-07-15 15:06:50 +00:00
|
|
|
num_grow_ch=32,
|
|
|
|
|
scale=4,
|
|
|
|
|
)
|
|
|
|
|
netscale = 4
|
2023-07-17 11:00:22 +00:00
|
|
|
elif self.model_name in ["RealESRGAN_x2plus.pth"]:
|
|
|
|
|
# x2 RRDBNet model
|
|
|
|
|
rrdbnet_model = RRDBNet(
|
|
|
|
|
num_in_ch=3,
|
|
|
|
|
num_out_ch=3,
|
|
|
|
|
num_feat=64,
|
|
|
|
|
num_block=23,
|
|
|
|
|
num_grow_ch=32,
|
|
|
|
|
scale=2,
|
|
|
|
|
)
|
|
|
|
|
netscale = 2
|
2023-07-15 15:06:50 +00:00
|
|
|
else:
|
|
|
|
|
msg = f"Invalid RealESRGAN model: {self.model_name}"
|
2024-01-13 12:23:16 +00:00
|
|
|
context.logger.error(msg)
|
2023-07-15 15:06:50 +00:00
|
|
|
raise ValueError(msg)
|
|
|
|
|
|
2024-04-12 04:55:21 +00:00
|
|
|
loadnet = context.models.load_ckpt_from_url(
|
|
|
|
|
source=ESRGAN_MODEL_URLS[self.model_name],
|
2024-03-19 01:11:13 +00:00
|
|
|
)
|
2023-07-15 15:06:50 +00:00
|
|
|
|
2024-04-25 03:53:30 +00:00
|
|
|
with loadnet as loadnet_model:
|
|
|
|
|
upscaler = RealESRGAN(
|
|
|
|
|
scale=netscale,
|
|
|
|
|
loadnet=loadnet_model,
|
|
|
|
|
model=rrdbnet_model,
|
|
|
|
|
half=False,
|
|
|
|
|
tile=self.tile_size,
|
|
|
|
|
)
|
2024-04-12 04:55:21 +00:00
|
|
|
|
2024-04-25 03:53:30 +00:00
|
|
|
# prepare image - Real-ESRGAN uses cv2 internally, and cv2 uses BGR vs RGB for PIL
|
|
|
|
|
# TODO: This strips the alpha... is that okay?
|
|
|
|
|
cv2_image = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
|
|
|
|
|
upscaled_image = upscaler.upscale(cv2_image)
|
2023-07-15 15:06:50 +00:00
|
|
|
|
2024-04-25 03:53:30 +00:00
|
|
|
pil_image = Image.fromarray(cv2.cvtColor(upscaled_image, cv2.COLOR_BGR2RGB)).convert("RGBA")
|
2023-10-04 04:23:31 +00:00
|
|
|
|
2024-01-13 12:23:16 +00:00
|
|
|
image_dto = context.images.save(image=pil_image)
|
Partial migration of UI to nodes API (#3195)
* feat(ui): add axios client generator and simple example
* fix(ui): update client & nodes test code w/ new Edge type
* chore(ui): organize generated files
* chore(ui): update .eslintignore, .prettierignore
* chore(ui): update openapi.json
* feat(backend): fixes for nodes/generator
* feat(ui): generate object args for api client
* feat(ui): more nodes api prototyping
* feat(ui): nodes cancel
* chore(ui): regenerate api client
* fix(ui): disable OG web server socket connection
* fix(ui): fix scrollbar styles typing and prop
just noticed the typo, and made the types stronger.
* feat(ui): add socketio types
* feat(ui): wip nodes
- extract api client method arg types instead of manually declaring them
- update example to display images
- general tidy up
* start building out node translations from frontend state and add notes about missing features
* use reference to sampler_name
* use reference to sampler_name
* add optional apiUrl prop
* feat(ui): start hooking up dynamic txt2img node generation, create middleware for session invocation
* feat(ui): write separate nodes socket layer, txt2img generating and rendering w single node
* feat(ui): img2img implementation
* feat(ui): get intermediate images working but types are stubbed out
* chore(ui): add support for package mode
* feat(ui): add nodes mode script
* feat(ui): handle random seeds
* fix(ui): fix middleware types
* feat(ui): add rtk action type guard
* feat(ui): disable NodeAPITest
This was polluting the network/socket logs.
* feat(ui): fix parameters panel border color
This commit should be elsewhere but I don't want to break my flow
* feat(ui): make thunk types more consistent
* feat(ui): add type guards for outputs
* feat(ui): load images on socket connect
Rudimentary
* chore(ui): bump redux-toolkit
* docs(ui): update readme
* chore(ui): regenerate api client
* chore(ui): add typescript as dev dependency
I am having trouble with TS versions after vscode updated and now uses TS 5. `madge` has installed 3.9.10 and for whatever reason my vscode wants to use that. Manually specifying 4.9.5 and then setting vscode to use that as the workspace TS fixes the issue.
* feat(ui): begin migrating gallery to nodes
Along the way, migrate to use RTK `createEntityAdapter` for gallery images, and separate `results` and `uploads` into separate slices. Much cleaner this way.
* feat(ui): clean up & comment results slice
* fix(ui): separate thunk for initial gallery load so it properly gets index 0
* feat(ui): POST upload working
* fix(ui): restore removed type
* feat(ui): patch api generation for headers access
* chore(ui): regenerate api
* feat(ui): wip gallery migration
* feat(ui): wip gallery migration
* chore(ui): regenerate api
* feat(ui): wip refactor socket events
* feat(ui): disable panels based on app props
* feat(ui): invert logic to be disabled
* disable panels when app mounts
* feat(ui): add support to disableTabs
* docs(ui): organise and update docs
* lang(ui): add toast strings
* feat(ui): wip events, comments, and general refactoring
* feat(ui): add optional token for auth
* feat(ui): export StatusIndicator and ModelSelect for header use
* feat(ui) working on making socket URL dynamic
* feat(ui): dynamic middleware loading
* feat(ui): prep for socket jwt
* feat(ui): migrate cancelation
also updated action names to be event-like instead of declaration-like
sorry, i was scattered and this commit has a lot of unrelated stuff in it.
* fix(ui): fix img2img type
* chore(ui): regenerate api client
* feat(ui): improve InvocationCompleteEvent types
* feat(ui): increase StatusIndicator font size
* fix(ui): fix middleware order for multi-node graphs
* feat(ui): add exampleGraphs object w/ iterations example
* feat(ui): generate iterations graph
* feat(ui): update ModelSelect for nodes API
* feat(ui): add hi-res functionality for txt2img generations
* feat(ui): "subscribe" to particular nodes
feels like a dirty hack but oh well it works
* feat(ui): first steps to node editor ui
* fix(ui): disable event subscription
it is not fully baked just yet
* feat(ui): wip node editor
* feat(ui): remove extraneous field types
* feat(ui): nodes before deleting stuff
* feat(ui): cleanup nodes ui stuff
* feat(ui): hook up nodes to redux
* fix(ui): fix handle
* fix(ui): add basic node edges & connection validation
* feat(ui): add connection validation styling
* feat(ui): increase edge width
* feat(ui): it blends
* feat(ui): wip model handling and graph topology validation
* feat(ui): validation connections w/ graphlib
* docs(ui): update nodes doc
* feat(ui): wip node editor
* chore(ui): rebuild api, update types
* add redux-dynamic-middlewares as a dependency
* feat(ui): add url host transformation
* feat(ui): handle already-connected fields
* feat(ui): rewrite SqliteItemStore in sqlalchemy
* fix(ui): fix sqlalchemy dynamic model instantiation
* feat(ui, nodes): metadata wip
* feat(ui, nodes): models
* feat(ui, nodes): more metadata wip
* feat(ui): wip range/iterate
* fix(nodes): fix sqlite typing
* feat(ui): export new type for invoke component
* tests(nodes): fix test instantiation of ImageField
* feat(nodes): fix LoadImageInvocation
* feat(nodes): add `title` ui hint
* feat(nodes): make ImageField attrs optional
* feat(ui): wip nodes etc
* feat(nodes): roll back sqlalchemy
* fix(nodes): partially address feedback
* fix(backend): roll back changes to pngwriter
* feat(nodes): wip address metadata feedback
* feat(nodes): add seeded rng to RandomRange
* feat(nodes): address feedback
* feat(nodes): move GET images error handling to DiskImageStorage
* feat(nodes): move GET images error handling to DiskImageStorage
* fix(nodes): fix image output schema customization
* feat(ui): img2img/txt2img -> linear
- remove txt2img and img2img tabs
- add linear tab
- add initial image selection to linear parameters accordion
* feat(ui): tidy graph builders
* feat(ui): tidy misc
* feat(ui): improve invocation union types
* feat(ui): wip metadata viewer recall
* feat(ui): move fonts to normal deps
* feat(nodes): fix broken upload
* feat(nodes): add metadata module + tests, thumbnails
- `MetadataModule` is stateless and needed in places where the `InvocationContext` is not available, so have not made it a `service`
- Handles loading/parsing/building metadata, and creating png info objects
- added tests for MetadataModule
- Lifted thumbnail stuff to util
* fix(nodes): revert change to RandomRangeInvocation
* feat(nodes): address feedback
- make metadata a service
- rip out pydantic validation, implement metadata parsing as simple functions
- update tests
- address other minor feedback items
* fix(nodes): fix other tests
* fix(nodes): add metadata service to cli
* fix(nodes): fix latents/image field parsing
* feat(nodes): customise LatentsField schema
* feat(nodes): move metadata parsing to frontend
* fix(nodes): fix metadata test
---------
Co-authored-by: maryhipp <maryhipp@gmail.com>
Co-authored-by: Mary Hipp <maryhipp@Marys-MacBook-Air.local>
2023-04-22 03:10:20 +00:00
|
|
|
|
2024-01-13 12:23:16 +00:00
|
|
|
return ImageOutput.build(image_dto)
|
