run on Mac Apple m1 pro

[alyfreym]

Hi, it is possible to run on MacBook M1 Pro?

[gtnbssn]

I see pytorch has mps support now. I would think it might be possible. Maybe by changing some relevant lines? For example in generate.py:

https://github.com/NVlabs/stylegan2-ada-pytorch/blob/d4b2afe9c27e3c305b721bc886d2cb5229458eba/generate.py#L82

would instead be:

    device = torch.device('mps')

But that's a very naive approach and I haven't tried. I don't even have a m1 or m2.

Has anyone tried?

[gtnbssn]

Just tried on a friend's m2.

Unfortunately this small change is not enough, as I get the following error:

TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

I tried changing this:

https://github.com/NVlabs/stylegan2-ada-pytorch/blob/d4b2afe9c27e3c305b721bc886d2cb5229458eba/generate.py#L119

To this:

        img = G(z, label, truncation_psi=truncation_psi, noise_mode=noise_mode, force_fp32=True)

(A trick that works if you want to generate images on the cpu: https://github.com/NVlabs/stylegan2-ada-pytorch/issues/105)

But I am still getting the same error.

[gtnbssn]

@alyfreym did you get it to work?

I'd argue this deserves to stay open, other people might be trying to achieve this, or come up with solutions.

I have done more digging and ran into more issues, so maybe the answer to your original questions is "not without converting the whole upfirdn2d code to something that runs on mps". It would still be nice to achieve it though!

If you don't want the notifications I think you can unsubscribe from the issue.

[NotNANtoN]

I got it to work. Update to newest MPS Pytorch version. pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu -U.

Then you do the changes above, making sure to also cast "z" to float before sending it to mps. I get 0.5 FPS with CPU and 2.9 FPS with the MPS on my MacBook with an M1.

[gtnbssn]

I got it to work. Update to newest MPS Pytorch version. pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu -U.

Then you do the changes above, making sure to also cast "z" to float before sending it to mps. I get 0.5 FPS with CPU and 2.9 FPS with the MPS on my MacBook with an M1.

Are you getting same results than on cuda though? I find that for the same parameters, my images look different, much darker. Would casting things to float have an influence on this?

Image generated with cuda on google colab:

With mps on my laptop:

[yuval-a]

Full steps:

1. Clone repo.
2. Install MiniForge
3. Create and activate Conda environment: conda create --name pytorch_env python=3.9 conda activate pytorch_env
4. Install PyTorch and Torch-Metal conda install pytorch torchvision torchaudio torch-metal -c pytorch
5. Replace all instances of torch.device(‘cuda’) in the code with torch.device(‘mps’)
6. In generate.py change to:

   # Create a numpy array of type float32
   random_array = np.random.RandomState(seed).randn(1, G.z_dim).astype(np.float32)
   # Convert the numpy array to a torch tensor and move it to the MPS device
   z = torch.from_numpy(random_array).to(device)