[Bug] Studio (macOS/MPS): image generation with an input image fails — VAE tiling builds float64 weights on MPS #12935

Environment

Where are you using Unsloth?

• Unsloth desktop application
• Unsloth web UI (unsloth studio)
• Unsloth CLI
• Python package or notebook
• Colab or Kaggle

Operating system and version: macOS 27.2

GPU model(s) and accelerator backend: Apple M4 Pro (48 GB unified memory), MPS (diffusers engine)

Versions:

• Unsloth Studio (Desktop): 0.1.903-beta
• Bundled backend: Python 3.13, torch 2.11.0, diffusers 0.41.0.dev0

What happened?

Steps to reproduce:

1. On an Apple Silicon Mac, open the Images tab and load a Qwen-Image-2.1 GGUF (Q4_K_M).
2. Start a generation that includes an input / reference image, so the pipeline VAE-encodes it.
3. The generation fails before the first denoise step.

Expected behavior: The input image is encoded and generation runs.

Actual behavior: Every generation with an input image fails with:

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

Model and operation involved: Diffusion inference, Qwen-Image-2.1 GGUF, VAE encode of the input image (tiled encode path).

Cause

axis_weights() in studio/backend/core/inference/diffusion_vae_tiling.py builds the tile blend weights as float64 directly on the target device:

pos = (torch.arange(size, dtype = torch.float64, device = device) + 0.5) / scale
total = torch.zeros(length * scale, dtype = torch.float64, device = device)
...
w = torch.ones(size, dtype = torch.float64, device = device)

MPS has no float64, so this raises as soon as tiled_encode runs on device=mps. The code arrived with #12696 and is unchanged by #12736 and #12764; it is still the same on main.

Suggested fix (tested locally)

Build the weights on the CPU in float64 and move only the fp32 result to device:

     size = min(tile, length) * scale
-    pos = (torch.arange(size, dtype = torch.float64, device = device) + 0.5) / scale
-    total = torch.zeros(length * scale, dtype = torch.float64, device = device)
+    pos = (torch.arange(size, dtype = torch.float64) + 0.5) / scale
+    total = torch.zeros(length * scale, dtype = torch.float64)
     weights = []
     for s in starts:
-        w = torch.ones(size, dtype = torch.float64, device = device)
+        w = torch.ones(size, dtype = torch.float64)
         ...
-    return [(w / total[s * scale : s * scale + size]).float() for w, s in zip(weights, starts)]
+    return [(w / total[s * scale : s * scale + size]).float().to(device) for w, s in zip(weights, starts)]

With this change, axis_weights([0, 48, 96], 64, 160, 8, torch, torch.device("mps")) returns fp32 tensors on mps:0 whose overlapping weights still sum to 1 at every pixel, and the error no longer occurs. The vectors are tiny (one per tile axis), so the CPU work and copy are negligible.

Diagnostics and logs

Traceback (trimmed, paths relative to site-packages/):

  File "diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py", line 463, in prepare_latents
    encoded = self._encode_vae_image(image, generator)
  File "diffusers/pipelines/qwenimage21/pipeline_qwenimage21.py", line 431, in _encode_vae_image
    image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax")
  File "studio/backend/core/inference/diffusion.py", line 9170, in _encode
    return original(x, *args, **kwargs)
  File "diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py", line 1297, in encode
    h = self._encode(x)
  File "diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py", line 1257, in _encode
    return self.tiled_encode(x)
  File "studio/backend/core/inference/diffusion_vae_tiling.py", line 587, in _tiled_encode
    return tiled_encode(self, x)
  File "studio/backend/core/inference/diffusion_vae_tiling.py", line 388, in tiled_encode
    wy = axis_weights(hs, tile, height, 1, torch, x.device, margin, ramp)
  File "studio/backend/core/inference/diffusion_vae_tiling.py", line 70, in axis_weights
    pos = (torch.arange(size, dtype = torch.float64, device = device) + 0.5) / scale
TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

Additional context

CUDA is unaffected because it supports float64; this only breaks on MPS (Apple Silicon).