offload_state_dict=True leaves CPU model weights on meta after loading #15018

UNet2DModel.from_pretrained(..., offload_state_dict=True) can return a model whose parameters remain on meta; its next forward pass fails. With device_map={"": "cpu"}, loading itself fails with “Cannot copy out of meta tensor”. Temporary weight files are also left behind. The same checkpoint loads and runs normally with offload_state_dict=False.

This reproduces on current main with a small, locally created model; no downloaded weights are needed:

import tempfile

import torch
from diffusers import UNet2DModel

model = UNet2DModel(
    sample_size=8,
    in_channels=3,
    out_channels=3,
    layers_per_block=1,
    block_out_channels=(4, 8),
    down_block_types=("DownBlock2D", "DownBlock2D"),
    up_block_types=("UpBlock2D", "UpBlock2D"),
    norm_num_groups=2,
)
with tempfile.TemporaryDirectory() as checkpoint:
    model.save_pretrained(checkpoint, max_shard_size="10KB")
    restored = UNet2DModel.from_pretrained(checkpoint, offload_state_dict=True)
    restored(torch.randn(1, 3, 8, 8), timestep=1)

Expected: temporary CPU-state offloading reduces loading memory without changing the resulting model's usability or weights.

Actual: RuntimeError: Tensor on device meta is not on the expected device cpu!

The restoration in _load_pretrained_model is nested under the condition for a nonempty permanent disk-offload index. Temporary CPU weight offloading does not require permanent disk placement, so that condition skips both restoration and cleanup.

I have prepared a local fix that restores a populated temporary index independently of the permanent index and always removes the temporary directory after successful loading. It also preserves the no-Accelerate fallback. Tests cover single-file/sharded checkpoints, serial/parallel loading, CPU placement, and CPU+disk placement with real inference comparisons. Would this scope be suitable for a PR?

Environment reproduced: Diffusers 0.41.0.dev0, PyTorch 2.14.1, Accelerate 1.15.0, Python 3.11.15, macOS arm64. The fix also passes the focused model-loading tests with PyTorch 2.6.0, Accelerate 0.31.0, and Python 3.10.20.