prevent accelerate from recreating keys such as and so that the VAE decode will progress. resolves #4526

2026-09-19 01:04:32 +02:00 · 2026-01-07 22:19:11 -05:00
parent 8d03d7c5b2
commit c1c4d6e0ee
1 changed files with 2 additions and 0 deletions
@@ -254,10 +254,12 @@ class OffloadHook(accelerate.hooks.ModelHook):
            if debug:
                shared.log.trace(f'Offload: type=balanced op=dispatch map={device_map}')
            if device_map is not None:
+                skip_keys = getattr(module, "_skip_keys", None)
                module = accelerate.dispatch_model(module,
                                                   main_device=torch.device(devices.device),
                                                   device_map=device_map,
                                                   offload_dir=offload_dir,
+                                                   skip_keys=skip_keys,
                                                   force_hooks=True,
                                                  )
            module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access