Spaces:
Runtime error
Runtime error
File size: 2,631 Bytes
6831a54 |
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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
import os
import sys
INITIALIZED = False
MONITOR_MODEL_MOVING = False
def monitor_module_moving():
if not MONITOR_MODEL_MOVING:
return
import torch
import traceback
old_to = torch.nn.Module.to
def new_to(*args, **kwargs):
traceback.print_stack()
print('Model Movement')
return old_to(*args, **kwargs)
torch.nn.Module.to = new_to
return
def initialize_forge():
global INITIALIZED
if INITIALIZED:
return
INITIALIZED = True
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(__file__)), 'packages_3rdparty'))
bad_list = ['--lowvram', '--medvram', '--medvram-sdxl']
for bad in bad_list:
if bad in sys.argv:
print(f'Arg {bad} is removed in Forge.')
print(f'Now memory management is fully automatic and you do not need any command flags.')
print(f'Please just remove this flag.')
print(f'In extreme cases, if you want to force previous lowvram/medvram behaviors, '
f'please use --always-offload-from-vram')
from backend.args import args
if args.gpu_device_id is not None:
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_device_id)
print("Set device to:", args.gpu_device_id)
if args.cuda_malloc:
from modules_forge.cuda_malloc import try_cuda_malloc
try_cuda_malloc()
from backend import memory_management
import torch
monitor_module_moving()
device = memory_management.get_torch_device()
torch.zeros((1, 1)).to(device, torch.float32)
memory_management.soft_empty_cache()
if memory_management.can_install_bnb():
from modules_forge.bnb_installer import try_install_bnb
try_install_bnb()
from backend import stream
print('CUDA Using Stream:', stream.should_use_stream())
from modules_forge.shared import diffusers_dir
# if 'TRANSFORMERS_CACHE' not in os.environ:
# os.environ['TRANSFORMERS_CACHE'] = diffusers_dir
if 'HF_HOME' not in os.environ:
os.environ['HF_HOME'] = diffusers_dir
if 'HF_DATASETS_CACHE' not in os.environ:
os.environ['HF_DATASETS_CACHE'] = diffusers_dir
if 'HUGGINGFACE_HUB_CACHE' not in os.environ:
os.environ['HUGGINGFACE_HUB_CACHE'] = diffusers_dir
if 'HUGGINGFACE_ASSETS_CACHE' not in os.environ:
os.environ['HUGGINGFACE_ASSETS_CACHE'] = diffusers_dir
if 'HF_HUB_CACHE' not in os.environ:
os.environ['HF_HUB_CACHE'] = diffusers_dir
import modules_forge.patch_basic
modules_forge.patch_basic.patch_all_basics()
return
|