Add simplified model manager install API to InvocationContext

[lstein]

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,...)

---

2. 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,...)

---

3. 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)

[lstein]

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.

[psychedelicious]

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.

[psychedelicious]

This function is unused - I think the logic to get the loader should be checking if it's a directory? I'm not sure how to fix this myself bc the diffusers_load_directory function has a different type signature than the other loader function options.

[lstein]

Fixed. Looks like I never wired up that function! In the process I discovered a long-standing bug that would prevent text encoders from being loaded generically.

I also added a unit test for loading generic diffusers models. To do this, I added the 10 MB "tiny" taesdxl model, which was the smallest loadable diffusers I could find.

[psychedelicious]

I don't think we should be adding models to the repo. Feels arbitrary to do it just for this one scenario and not all of them, plus 10MB increases the repo size by like 8%. Don't wanna get in the habit of doing this.

I think the best way to test model loading is to package up a models dir that has samples of various model types and configurations, then run full loading tests locally on one of our own machines or a custom CI runner.

[lstein]

I've commented out the unit test for directory loading.

[psychedelicious]

Can you please drop the commits that added the model file so it's not in the repo?

[psychedelicious]

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

Two minor issues noted.

[lstein]

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

[lstein]

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]

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]

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

[lstein]

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]

@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]

@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.

[RyanJDick]

I just reviewed the invocation_context.py API.

[RyanJDick]

This docstring is outdated.

[lstein]

Fixed.

[RyanJDick]

I don't love that we switch behaviour based on whether source is a Path or a str. It feels like a fragile distinction, especially given the popularity of using strs to represent paths in Python. The caller should always know whether they are dealing with a path or a URL/repo name, so I think it's better to make this distinction explicit.

In this discussion we had landed on an API that didn't require this type condition. Was there a reason for moving away from that?

[lstein]

It looks like @psychedelicious removed load_model_from_url() in commit b12444 . Add it back?

[psychedelicious]

I moved the method from the high-level services.model_manager class to the invocation context. The function is otherwise the same.

[lstein]

@psychedelicious I can’t find load_model_from_url() in invocation_context.py, and the commit indicates that it was deleted from the model manager and then the load_and_cache_model() was modified to no longer call it.

I’ll take care of fixing the API if the change wasn’t deliberate.

[psychedelicious]

load_model_from_url() was this:

class ModelManagerService(ModelManagerServiceBase):
    # ...
    def load_model_from_url(
        self,
        source: str | AnyHttpUrl,
        loader: Optional[Callable[[Path], Dict[str, torch.Tensor]]] = None,
    ) -> LoadedModel:
        model_path = self.install.download_and_cache_model(source=str(source))
        return self.load.load_model_from_path(model_path=model_path, loader=loader)

It was called in one place in invocation_context.py:

class ModelsInterface(InvocationContextInterface):
    def load_and_cache_model(
        self,
        source: Path | str | AnyHttpUrl,
        loader: Optional[Callable[[Path], dict[str, Tensor]]] = None,
    ) -> LoadedModel:
        if isinstance(source, Path):
            return self._services.model_manager.load.load_model_from_path(model_path=source, loader=loader)
        else:
            # Called here
            return self._services.model_manager.load_model_from_url(source=source, loader=loader)

What I did was copy the two lines from the function body directly into ModelsInterface.load_and_cache_model() - I don't think we should have load_model_from_url on the main ModelManagerService.

I think the API Ryan is suggesting is to have load_and_cache_from_path and load_and_cache_from_url on ModelsInterface

[lstein]

For clarity, I have renamed the methods load_local_model() and load_remote_model(). The former accepts the Path to a file or directory, and the latter accepts either a direct download URL or a HuggingFace URL. I have fixed the documentation and updated the pull request description.

[psychedelicious]

Sounds good!

[lstein]

@RyanJDick I've fixed the issues you identified.

[RyanJDick]

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.