update

app.py

@@ -31,7 +31,7 @@ def load_models(config_name, ckpt_path, vae_path, device):
 
     # Load main U-Net model
     unet = MaskDiT(**params['model']).to(device)
-    unet.load_state_dict(torch.load(ckpt_path)['model'])
+    unet.load_state_dict(torch.load(ckpt_path, map_location='cpu')['model'])
     unet.eval()
 
     accelerator = Accelerator(mixed_precision="fp16")