Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add simplified model manager install API to InvocationContext #6132

Merged
merged 60 commits into from
Jun 8, 2024

Conversation

lstein
Copy link
Collaborator

@lstein lstein commented Apr 4, 2024

Summary

This three two model manager-related methods to the InvocationContext uniform API. They are accessible via context.models.*:

  1. load_local_model(model_path: Path, loader: Optional[Callable[[Path], AnyModel]] = None) -> LoadedModelWithoutConfig

Load the model located at the indicated path.

This will load a local model (.safetensors, .ckpt or diffusers directory) into the model manager RAM cache and return its LoadedModelWithoutConfig. If the optional loader argument is provided, the loader will be invoked to load the model into memory. Otherwise the method will call safetensors.torch.load_file() torch.load() (with a pickle scan), or from_pretrained() as appropriate to the path type.

Be aware that the LoadedModelWithoutConfig object differs from LoadedModel by having no config attribute.

Here is an example of usage:

def invoke(self, context: InvocatinContext) -> ImageOutput:
       model_path = Path('/opt/models/RealESRGAN_x4plus.pth')
       loadnet = context.models.load_local_model(model_path)
       with loadnet as loadnet_model:
             upscaler = RealESRGAN(loadnet=loadnet_model,...)

  1. load_remote_model(source: str | AnyHttpUrl, loader: Optional[Callable[[Path], AnyModel]] = None) -> LoadedModelWithoutConfig

Load the model located at the indicated URL or repo_id.

This is similar to load_local_model() but it accepts either a HugginFace repo_id (as a string), or a URL. The model's file(s) will be downloaded to models/.download_cache and then loaded, returning a

def invoke(self, context: InvocatinContext) -> ImageOutput:
       model_url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'
       loadnet = context.models.load_remote_model(model_url)
       with loadnet as loadnet_model:
             upscaler = RealESRGAN(loadnet=loadnet_model,...)

  1. download_and_cache_model( source: str | AnyHttpUrl, access_token: Optional[str] = None, timeout: Optional[int] = 0) -> Path

Download the model file located at source to the models cache and return its Path. This will check models/.download_cache for the desired model file and download it from the indicated source if not already present. The local Path to the downloaded file is then returned.


Other Changes

This PR performs a migration, in which it renames models/.cache to models/.convert_cache, and migrates previously-downloaded ESRGAN, openpose, DepthAnything and Lama inpaint models from the models/core directory into models/.download_cache.

There are a number of legacy model files in models/core, such as GFPGAN, which are no longer used. This PR deletes them and tidies up the models/core directory.

Related Issues / Discussions

I have systematically replaced all the calls to download_with_progress_bar(). This function is no longer used elsewhere and has been removed.

QA Instructions

I have added unit tests for the three new calls. You may test that the load_and_cache_model() call is working by running the upscaler within the web app. On first try, you will see the model file being downloaded into the models .cache directory. On subsequent tries, the model will either load from RAM (if it hasn't been displaced) or will be loaded from the filesystem.

Merge Plan

Squash merge when approved.

Checklist

  • The PR has a short but descriptive title, suitable for a changelog
  • Tests added / updated (if applicable)
  • Documentation added / updated (if applicable)

@github-actions github-actions bot added python PRs that change python files services PRs that change app services labels Apr 4, 2024
@lstein lstein force-pushed the lstein/feat/simple-mm2-api branch from 9cc1f20 to af1b57a Compare April 12, 2024 01:46
@github-actions github-actions bot added invocations PRs that change invocations backend PRs that change backend files python-tests PRs that change python tests labels Apr 12, 2024
@lstein lstein marked this pull request as ready for review April 12, 2024 05:17
@lstein
Copy link
Collaborator Author

lstein commented Apr 14, 2024

I have added a migration script that tidies up the models/core directory and removes unused models such as GFPGAN. In addition, I have renamed models/.cache to models/.convert_cache to distinguish it from the directory in which just-in-time models are downloaded into, which is models/.download_cache. While the size of models/.convert_cache is capped such that less-used models are cleared periodically, files in models/.download_cache are not removed unless the user does so manually.

@lstein lstein force-pushed the lstein/feat/simple-mm2-api branch from 537a626 to 3ddd7ce Compare April 14, 2024 19:57
@lstein lstein force-pushed the lstein/feat/simple-mm2-api branch from 3ddd7ce to fa6efac Compare April 14, 2024 20:10
Copy link
Collaborator

@psychedelicious psychedelicious left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure what I was expecting the implementation to be, but it definitely wasn't as simple as this - great work.

I've requested a few changes and there's one discussion item that I'd like to marinate on before we change the public invocation API.

invokeai/app/invocations/upscale.py Outdated Show resolved Hide resolved
invokeai/app/services/shared/invocation_context.py Outdated Show resolved Hide resolved
@psychedelicious psychedelicious self-requested a review June 3, 2024 01:54
Copy link
Collaborator

@psychedelicious psychedelicious left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry for the delay in reviewing. I've tidied a few things and tested everything, working great!

Two minor issues noted.

@lstein
Copy link
Collaborator Author

lstein commented Jun 4, 2024

@psychedelicious I've addressed the remaining issues you raised. Thanks for a thorough review.

@lstein
Copy link
Collaborator Author

lstein commented Jun 4, 2024

I removed a number of unnecessary changes in invocation_context.py, mostly extraneous type annotations. If mypy is complaining about these, then that's a mypy problem, because all the methods are annotated correctly.

I also moved load_model_from_url from the main model manager class into the invocation context.

Yes, mypy is having trouble tracking the return type of several methods. I haven't figured out what causes the problem and don't want to add a # type: ignore. But maybe I should 'cause I'm not ready to turn to pyright.

@psychedelicious
Copy link
Collaborator

Yes, mypy is having trouble tracking the return type of several methods. I haven't figured out what causes the problem and don't want to add a # type: ignore. But maybe I should 'cause I'm not ready to turn to pyright.

We shouldn't add # type: ignore, that will stop all type checkers from doing anything - including pyright. The places where you made code quality concessions to satisfy mypy involve very straightforward types - either your mypy config is borked or mypy itself is borked. FWIW, I've found pyright to be much faster, more thorough and more correct than mypy.

@psychedelicious
Copy link
Collaborator

@RyanJDick Would you mind doing one last review of this PR?

@lstein
Copy link
Collaborator Author

lstein commented Jun 4, 2024

Yes, mypy is having trouble tracking the return type of several methods. I haven't figured out what causes the problem and don't want to add a # type: ignore. But maybe I should 'cause I'm not ready to turn to pyright.

We shouldn't add # type: ignore, that will stop all type checkers from doing anything - including pyright. The places where you made code quality concessions to satisfy mypy involve very straightforward types - either your mypy config is borked or mypy itself is borked. FWIW, I've found pyright to be much faster, more thorough and more correct than mypy.

You've convinced me. I've switched to pyright!

@RyanJDick
Copy link
Collaborator

@RyanJDick Would you mind doing one last review of this PR?

Looks like 43/44 files have changed since I last looked 😅 . I'll plan to spend a chunk of time on this tomorrow.

@psychedelicious
Copy link
Collaborator

@RyanJDick Can narrow that down to reviewing invocation_context.py, which changes the public API and is more important to get right the first time. Thanks.

Copy link
Collaborator

@RyanJDick RyanJDick left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I just reviewed the invocation_context.py API.

@lstein lstein requested a review from RyanJDick June 6, 2024 02:57
@lstein
Copy link
Collaborator Author

lstein commented Jun 6, 2024

@RyanJDick I've fixed the issues you identified.

Copy link
Collaborator

@RyanJDick RyanJDick left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

invocation_context.py looks good to me ✅

I'll defer to @psychedelicious for final approval, since he has a more complete understanding of this PR than me at this point.

@github-actions github-actions bot added the frontend PRs that change frontend files label Jun 7, 2024
@hipsterusername hipsterusername merged commit 9432336 into main Jun 8, 2024
14 checks passed
@hipsterusername hipsterusername deleted the lstein/feat/simple-mm2-api branch June 8, 2024 23:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
api backend PRs that change backend files docs PRs that change docs frontend PRs that change frontend files invocations PRs that change invocations python PRs that change python files python-tests PRs that change python tests services PRs that change app services
Projects
None yet
Development

Successfully merging this pull request may close these issues.

[bug]: CUDA out of memory error when upscaling x4 (or x2 twice)
4 participants