[FEATURE] Improve overall device placement handling

[davidberenstein1957]

Is your feature request related to a problem? Please describe.

I noticed this broader habbit in #97 and partially covered there, but I realised there is a wider habit of enforcing cuda usage, although people might still be interested in using cpu or mps, and @sharpenb mentioned we wanted to try supporting it.

Describe the solution you'd like

I propose re-using the logic in #97 and doing some other things:

• Always check and warm for device placement, where we prefer mps/cuda when no preference is given.
• Warn when using mps about limited support while nudging users to open an issue in case of weird behaviour. (torch..compile implements internal MPS warning, for example)

Additional context

[davidberenstein1957]

WDYT. @gsprochette @sharpenb @johannaSommer

[davidberenstein1957]

This also seems to be the case within Evaluation Metrics as I figured out together with @begumcig .

For example, ElapsedTimeMetric automatically usestorch.cuda.Event

[gsprochette]

I completely agree with your two points. Since now the device in SmashConfig should be correctly handled we should really use this one as much as possible. The evaluation doesn't have one so it needs an update with the new device utils (probably changing the default of Task's device argument and using task.device whenever a device is need), but other than this it should be straightforward, no?

[davidberenstein1957]

@gsprochette, yes, it should be pretty straightforward. I mostly wanted to align thoughts on this before starting any work.

[gsprochette]

Cool, just wanted to make sure of the scope :)

[davidberenstein1957]

I just had a similar experience when loading a SanaPipeline

The following does not work.

from pruna import PrunaModel

optimized_model = PrunaModel.from_hub("davidberenstein1957/Sana_600M_512px_diffusers")

The following works, but loads part of the architecture on different devices.

import torch

from pruna import PrunaModel

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
from pruna import PrunaModel

optimized_model = PrunaModel.from_hub("davidberenstein1957/Sana_600M_512px_diffusers", device=device)

# Run inference
optimized_model("A serene landscape with mountains").images[0]

The following is a working load pipeline.

import torch

from pruna import PrunaModel

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

optimized_model = PrunaModel.from_hub("davidberenstein1957/Sana_600M_512px_diffusers", device=device)
optimized_model.model.transformer.to_empty(device=device)  # Set transformer to default dtype from model card
optimized_model.model.transformer.to(torch.float16)  # Set text encoder to bfloat16
optimized_model.model.text_encoder.to(device, torch.bfloat16)  # Set text encoder to bfloat16
optimized_model.model.vae.to(device, torch.float32)  # Set VAE to float32 for better stability

# Run inference
optimized_model("A serene landscape with mountains").images[0]