mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 01:04:32 +02:00
revert sd_models.py
This commit is contained in:
@@ -405,7 +405,7 @@ def apply_balanced_offload(sd_model):
|
||||
if hasattr(pipe, "_internal_dict"):
|
||||
keys = pipe._internal_dict.keys() # pylint: disable=protected-access
|
||||
else:
|
||||
keys = get_signature(pipe).keys()
|
||||
keys = get_signature(shared.sd_model).keys()
|
||||
for module_name in keys: # pylint: disable=protected-access
|
||||
module = getattr(pipe, module_name, None)
|
||||
if isinstance(module, torch.nn.Module):
|
||||
@@ -1448,14 +1448,10 @@ def disable_offload(sd_model):
|
||||
from accelerate.hooks import remove_hook_from_module
|
||||
if not getattr(sd_model, 'has_accelerate', False):
|
||||
return
|
||||
if hasattr(sd_model, "_internal_dict"):
|
||||
keys = sd_model._internal_dict.keys() # pylint: disable=protected-access
|
||||
else:
|
||||
keys = get_signature(sd_model).keys()
|
||||
for module_name in keys: # pylint: disable=protected-access
|
||||
module = getattr(sd_model, module_name, None)
|
||||
if isinstance(module, torch.nn.Module):
|
||||
module = remove_hook_from_module(module, recurse=True)
|
||||
if hasattr(sd_model, 'components'):
|
||||
for _name, model in sd_model.components.items():
|
||||
if isinstance(model, torch.nn.Module):
|
||||
remove_hook_from_module(model, recurse=True)
|
||||
sd_model.has_accelerate = False
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user