File size: 1,677 Bytes
c144dd2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
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
unet_enabled = True
vae_enabled = True
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
|