[Diffusion] Add loader-owned host-weight plans for DLO
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
1 parent
d1e230c95b
commit
e2796b6168
17 files changed
+1717
-588
No files matched your search
@@ -292,7 +292,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
|
||||
)
|
||||
|
||||
# Apply CPU offloading
|
||||
self.offload_backend = get_offload_backend(self.od_config, device=self.device)
|
||||
self.offload_backend = get_offload_backend(
|
||||
self.od_config,
|
||||
device=self.device,
|
||||
host_weight_plan=model_loader.take_host_weight_plan(),
|
||||
)
|
||||
if self.offload_backend is not None:
|
||||
logger.info(f" Enabling offloader backend: {self.offload_backend.__class__.__name__}")
|
||||
self.offload_backend.enable(self.pipeline)
|
||||
|
||||
Reference in new issue
Block a user