[util] Add generic torch device class

[lstein]

Summary

This PR cleans up the various calls to choose_torch_device(), torch_dtype(), choose_precision() and so on, removes redundant function calls, and creates a single class named TorchDeviceSelect that supersedes their functionality. In addition to creating a simplified API, this class generalizes the call to clear the VRAM cache so that the same method empties the VRAM cache for both CUDA and MPS devices. It also provides an API for invocation context-dependent retrieval of the GPU device, intended for use in GPU load balancing in the future.

Example usage:

from invokeai.backend.util.devices import TorchDevice

class Foo(BaseInvocation):
    def invoke(self, context: InvocationContext):
         torch_device = TorchDevice.choose_torch_device()
         torch_dtype = TorchDevice.choose_torch_dtype()

         # empty CUDA or MPS cache
         TorchDevice.empty_cache()

The methods that return strings instead of objects, e.g. "float32" rather than torch.float32, have been removed. The legacy calls to choose_torch_device() and choose_precision() functions now issue a deprecation warning.

Related Issues / Discussions

QA Instructions

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)

[psychedelicious]

What's the difference between "auto" and "autocast"?

[lstein]

What's the difference between "auto" and "autocast"?

"auto" selects one of the floating point precision types. "autocast" used to activate the torch.autocast() context for certain generation operations, but I this code has been removed. So I've gone ahead and removed references to this configuration option.

By the way, I notice a TODO from ryan in model_patcher.py saying that torch.autocast might provide a speed benefit in model patching. It could be hard-coded there if needed.

[psychedelicious]

"auto" selects one of the floating point precision types. "autocast" used to activate the torch.autocast() context for certain generation operations, but I this code has been removed. So I've gone ahead and removed references to this configuration option.

Gotcha. We need to be careful with changing valid config settings. If somebody has "autocast" in their invokeai.yaml file, they will get a pydantic error on startup. It may be OK for this particular setting, but generally we'll want to write a config script migration, like config_default.py:migrate_v3_config_dict

[lstein]

Right. I'll put in a migrate script

[RyanJDick]

Good work on this PR. I love to see our test coverage going up! 😄

I've approved, but would still like some clarity on how you're thinking about the following:
1. As @psychedelicious mentioned this will break configs that contain autocast. Are we thinking that there are probably very few people with that setting? So we'll just provide guidance to anyone who does hit this?
Edit: I just saw your comment about addressing this.
2. How confident are we that context.models.get_execution_device() is the right API? I can imagine how it would enable multi-gpu - but without the full context of how it's going to be used I'm slightly nervous that we're going to end up making breaking changes to it.

[psychedelicious]

Indeed we should be careful when adding to the public API, because we are promising to support and maintain it.

Some more thoughts:

• I'm not sure get_execution_device should be in the context.models namespace. Feels more like a utility to me.
• There's no indication that using TorchDeviceSelect.get_execution_device() directly in a node is going to cause problems - because it doesn't except in a separate branch that may be merged in the future and not work for some nodes. In fact, some nodes do use TorchDeviceSelect directly. It's a confusing and has footgun potential.
• "TorchDeviceSelect" suggests to me that I can use this class to select a torch device. Maybe something more generic like "TorchHelper" is clearer.
• The invocation_api module exports some objects that are removed in this PR:

  from invokeai.backend.util.devices import CPU_DEVICE, CUDA_DEVICE, MPS_DEVICE, choose_precision, choose_torch_device

  For better or worse, we have committed these to the public API. Need to handle that.

[lstein]

I'm tired of working on this and will get back to it at some point soon. Converting to draft for now.

[lstein]

Indeed we should be careful when adding to the public API, because we are promising to support and maintain it.

Some more thoughts:

• I'm not sure get_execution_device should be in the context.models namespace. Feels more like a utility to me.

• There's no indication that using TorchDeviceSelect.get_execution_device() directly in a node is going to cause problems - because it doesn't except in a separate branch that may be merged in the future and not work for some nodes. In fact, some nodes do use TorchDeviceSelect directly. It's a confusing and has footgun potential.

• "TorchDeviceSelect" suggests to me that I can use this class to select a torch device. Maybe something more generic like "TorchHelper" is clearer.

• The invocation_api module exports some objects that are removed in this PR:

  from invokeai.backend.util.devices import CPU_DEVICE, CUDA_DEVICE, MPS_DEVICE, choose_precision, choose_torch_device

  For better or worse, we have committed these to the public API. Need to handle that.

These recommendations have now been implemented.

[psychedelicious]

Thanks for addressing the feedback, sorry if it was a hassle. Couple minor changes and comments