Add build-in extension "NeverOOM"

see also discussions
This commit is contained in:
lllyasviel
2024-02-24 19:09:06 -08:00
parent 50229a05c1
commit 437c348926
5 changed files with 68 additions and 6 deletions

View File

@@ -0,0 +1,47 @@
import gradio as gr
from modules import scripts
from ldm_patched.modules import model_management
class NeverOOMForForge(scripts.Script):
sorting_priority = 18
def __init__(self):
self.previous_unet_enabled = False
self.original_vram_state = model_management.vram_state
def title(self):
return "Never OOM Integrated"
def show(self, is_img2img):
return scripts.AlwaysVisible
def ui(self, *args, **kwargs):
with gr.Accordion(open=False, label=self.title()):
unet_enabled = gr.Checkbox(label='Enabled for UNet (always maximize offload)', value=False)
vae_enabled = gr.Checkbox(label='Enabled for VAE (always tiled)', value=False)
return unet_enabled, vae_enabled
def process(self, p, *script_args, **kwargs):
unet_enabled, vae_enabled = script_args
if unet_enabled:
print('NeverOOM Enabled for UNet (always maximize offload)')
if vae_enabled:
print('NeverOOM Enabled for VAE (always tiled)')
model_management.VAE_ALWAYS_TILED = vae_enabled
if self.previous_unet_enabled != unet_enabled:
model_management.unload_all_models()
if unet_enabled:
self.original_vram_state = model_management.vram_state
model_management.vram_state = model_management.VRAMState.NO_VRAM
else:
model_management.vram_state = self.original_vram_state
print(f'VARM State Changed To {model_management.vram_state.name}')
self.previous_unet_enabled = unet_enabled
return