fluxdev/stable-diffusion-webui-forge

1import os2import sys3 4 5MONITOR_MODEL_MOVING = False6 7 8def monitor_module_moving():9    if not MONITOR_MODEL_MOVING:10        return11 12    import torch13    import traceback14 15    old_to = torch.nn.Module.to16 17    def new_to(*args, **kwargs):18        traceback.print_stack()19        print('Model Movement')20 21        return old_to(*args, **kwargs)22 23    torch.nn.Module.to = new_to24    return25 26 27def initialize_forge():28    bad_list = ['--lowvram', '--medvram', '--medvram-sdxl']29 30    for bad in bad_list:31        if bad in sys.argv:32            print(f'Arg {bad} is removed in Forge.')33            print(f'Now memory management is fully automatic and you do not need any command flags.')34            print(f'Please just remove this flag.')35            print(f'In extreme cases, if you want to force previous lowvram/medvram behaviors, '36                  f'please use --always-offload-from-vram')37 38    from ldm_patched.modules import args_parser39 40    args_parser.args, _ = args_parser.parser.parse_known_args()41 42    if args_parser.args.gpu_device_id is not None:43        os.environ['CUDA_VISIBLE_DEVICES'] = str(args_parser.args.gpu_device_id)44        print("Set device to:", args_parser.args.gpu_device_id)45 46    if args_parser.args.cuda_malloc:47        from modules_forge.cuda_malloc import try_cuda_malloc48        try_cuda_malloc()49 50    import ldm_patched.modules.model_management as model_management51    import torch52 53    monitor_module_moving()54 55    device = model_management.get_torch_device()56    torch.zeros((1, 1)).to(device, torch.float32)57    model_management.soft_empty_cache()58 59    import modules_forge.patch_basic60    modules_forge.patch_basic.patch_all_basics()61 62    from modules_forge import stream63    print('CUDA Stream Activated: ', stream.using_stream)64 65    from modules_forge.shared import diffusers_dir66 67    if 'TRANSFORMERS_CACHE' not in os.environ:68        os.environ['TRANSFORMERS_CACHE'] = diffusers_dir69 70    if 'HF_HOME' not in os.environ:71        os.environ['HF_HOME'] = diffusers_dir72 73    if 'HF_DATASETS_CACHE' not in os.environ:74        os.environ['HF_DATASETS_CACHE'] = diffusers_dir75 76    if 'HUGGINGFACE_HUB_CACHE' not in os.environ:77        os.environ['HUGGINGFACE_HUB_CACHE'] = diffusers_dir78 79    if 'HUGGINGFACE_ASSETS_CACHE' not in os.environ:80        os.environ['HUGGINGFACE_ASSETS_CACHE'] = diffusers_dir81 82    if 'HF_HUB_CACHE' not in os.environ:83        os.environ['HF_HUB_CACHE'] = diffusers_dir84    return85