seperates out arg parser and imports args
This commit is contained in:
EllangoK
parent
dd29966f8a
commit
e5e587b1c0
|
|
@ -0,0 +1,29 @@
|
|||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument("--listen", type=str, default="127.0.0.1", help="Listen on IP or 127.0.0.1 if none given so the UI can be accessed from other computers.")
|
||||
parser.add_argument("--port", type=int, default=8188, help="Set the listen port.")
|
||||
parser.add_argument("--extra-model-paths-config", type=str, default=None, help="Load an extra_model_paths.yaml file.")
|
||||
parser.add_argument("--output-directory", type=str, default=None, help="Set the ComfyUI output directory.")
|
||||
parser.add_argument("--dont-upcast-attention", action="store_true", help="Disable upcasting of attention. Can boost speed but increase the chances of black images.")
|
||||
|
||||
attn_group = parser.add_mutually_exclusive_group()
|
||||
attn_group.add_argument("--use-split-cross-attention", action="store_true", help="Use the split cross attention optimization instead of the sub-quadratic one. Ignored when xformers is used.")
|
||||
attn_group.add_argument("--use-pytorch-cross-attention", action="store_true", help="Use the new pytorch 2.0 cross attention function.")
|
||||
|
||||
parser.add_argument("--disable-xformers", action="store_true", help="Disable xformers.")
|
||||
parser.add_argument("--cuda-device", type=int, default=None, help="Set the id of the cuda device this instance will use.")
|
||||
|
||||
vram_group = parser.add_mutually_exclusive_group()
|
||||
vram_group.add_argument("--highvram", action="store_true", help="By default models will be unloaded to CPU memory after being used. This option keeps them in GPU memory.")
|
||||
vram_group.add_argument("--normalvram", action="store_true", help="Used to force normal vram use if lowvram gets automatically enabled.")
|
||||
vram_group.add_argument("--lowvram", action="store_true", help="Split the unet in parts to use less vram.")
|
||||
vram_group.add_argument("--novram", action="store_true", help="When lowvram isn't enough.")
|
||||
vram_group.add_argument("--cpu", action="store_true", help="To use the CPU for everything (slow).")
|
||||
|
||||
parser.add_argument("--dont-print-server", action="store_true", help="Don't print server output.")
|
||||
parser.add_argument("--quick-test-for-ci", action="store_true", help="Quick test for CI.")
|
||||
parser.add_argument("--windows-standalone-build", action="store_true", help="Windows standalone build.")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
|
@ -21,6 +21,8 @@ if model_management.xformers_enabled():
|
|||
import os
|
||||
_ATTN_PRECISION = os.environ.get("ATTN_PRECISION", "fp32")
|
||||
|
||||
from cli_args import args
|
||||
|
||||
def exists(val):
|
||||
return val is not None
|
||||
|
||||
|
|
@ -474,7 +476,6 @@ class CrossAttentionPytorch(nn.Module):
|
|||
|
||||
return self.to_out(out)
|
||||
|
||||
import sys
|
||||
if model_management.xformers_enabled():
|
||||
print("Using xformers cross attention")
|
||||
CrossAttention = MemoryEfficientCrossAttention
|
||||
|
|
@ -482,7 +483,7 @@ elif model_management.pytorch_attention_enabled():
|
|||
print("Using pytorch cross attention")
|
||||
CrossAttention = CrossAttentionPytorch
|
||||
else:
|
||||
if "--use-split-cross-attention" in sys.argv:
|
||||
if args.use_split_cross_attention:
|
||||
print("Using split optimization for cross attention")
|
||||
CrossAttention = CrossAttentionDoggettx
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -1,36 +1,35 @@
|
|||
import psutil
|
||||
from enum import Enum
|
||||
from cli_args import args
|
||||
|
||||
CPU = 0
|
||||
NO_VRAM = 1
|
||||
LOW_VRAM = 2
|
||||
NORMAL_VRAM = 3
|
||||
HIGH_VRAM = 4
|
||||
MPS = 5
|
||||
class VRAMState(Enum):
|
||||
CPU = 0
|
||||
NO_VRAM = 1
|
||||
LOW_VRAM = 2
|
||||
NORMAL_VRAM = 3
|
||||
HIGH_VRAM = 4
|
||||
MPS = 5
|
||||
|
||||
accelerate_enabled = False
|
||||
vram_state = NORMAL_VRAM
|
||||
# Determine VRAM State
|
||||
vram_state = VRAMState.NORMAL_VRAM
|
||||
set_vram_to = VRAMState.NORMAL_VRAM
|
||||
|
||||
total_vram = 0
|
||||
total_vram_available_mb = -1
|
||||
|
||||
import sys
|
||||
import psutil
|
||||
|
||||
forced_cpu = "--cpu" in sys.argv
|
||||
|
||||
set_vram_to = NORMAL_VRAM
|
||||
accelerate_enabled = False
|
||||
|
||||
try:
|
||||
import torch
|
||||
total_vram = torch.cuda.mem_get_info(torch.cuda.current_device())[1] / (1024 * 1024)
|
||||
total_ram = psutil.virtual_memory().total / (1024 * 1024)
|
||||
forced_normal_vram = "--normalvram" in sys.argv
|
||||
if not forced_normal_vram and not forced_cpu:
|
||||
if not args.normalvram and not args.cpu:
|
||||
if total_vram <= 4096:
|
||||
print("Trying to enable lowvram mode because your GPU seems to have 4GB or less. If you don't want this use: --normalvram")
|
||||
set_vram_to = LOW_VRAM
|
||||
set_vram_to = VRAMState.LOW_VRAM
|
||||
elif total_vram > total_ram * 1.1 and total_vram > 14336:
|
||||
print("Enabling highvram mode because your GPU has more vram than your computer has ram. If you don't want this use: --normalvram")
|
||||
vram_state = HIGH_VRAM
|
||||
vram_state = VRAMState.HIGH_VRAM
|
||||
except:
|
||||
pass
|
||||
|
||||
|
|
@ -39,34 +38,32 @@ try:
|
|||
except:
|
||||
OOM_EXCEPTION = Exception
|
||||
|
||||
if "--disable-xformers" in sys.argv:
|
||||
XFORMERS_IS_AVAILBLE = False
|
||||
if args.disable_xformers:
|
||||
XFORMERS_IS_AVAILABLE = False
|
||||
else:
|
||||
try:
|
||||
import xformers
|
||||
import xformers.ops
|
||||
XFORMERS_IS_AVAILBLE = True
|
||||
XFORMERS_IS_AVAILABLE = True
|
||||
except:
|
||||
XFORMERS_IS_AVAILBLE = False
|
||||
XFORMERS_IS_AVAILABLE = False
|
||||
|
||||
ENABLE_PYTORCH_ATTENTION = False
|
||||
if "--use-pytorch-cross-attention" in sys.argv:
|
||||
ENABLE_PYTORCH_ATTENTION = args.use_pytorch_cross_attention
|
||||
if ENABLE_PYTORCH_ATTENTION:
|
||||
torch.backends.cuda.enable_math_sdp(True)
|
||||
torch.backends.cuda.enable_flash_sdp(True)
|
||||
torch.backends.cuda.enable_mem_efficient_sdp(True)
|
||||
ENABLE_PYTORCH_ATTENTION = True
|
||||
XFORMERS_IS_AVAILBLE = False
|
||||
XFORMERS_IS_AVAILABLE = False
|
||||
|
||||
if args.lowvram:
|
||||
set_vram_to = VRAMState.LOW_VRAM
|
||||
elif args.novram:
|
||||
set_vram_to = VRAMState.NO_VRAM
|
||||
elif args.highvram:
|
||||
vram_state = VRAMState.HIGH_VRAM
|
||||
|
||||
|
||||
if "--lowvram" in sys.argv:
|
||||
set_vram_to = LOW_VRAM
|
||||
if "--novram" in sys.argv:
|
||||
set_vram_to = NO_VRAM
|
||||
if "--highvram" in sys.argv:
|
||||
vram_state = HIGH_VRAM
|
||||
|
||||
|
||||
if set_vram_to == LOW_VRAM or set_vram_to == NO_VRAM:
|
||||
if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM):
|
||||
try:
|
||||
import accelerate
|
||||
accelerate_enabled = True
|
||||
|
|
@ -81,14 +78,14 @@ if set_vram_to == LOW_VRAM or set_vram_to == NO_VRAM:
|
|||
|
||||
try:
|
||||
if torch.backends.mps.is_available():
|
||||
vram_state = MPS
|
||||
vram_state = VRAMState.MPS
|
||||
except:
|
||||
pass
|
||||
|
||||
if forced_cpu:
|
||||
vram_state = CPU
|
||||
if args.cpu:
|
||||
vram_state = VRAMState.CPU
|
||||
|
||||
print("Set vram state to:", ["CPU", "NO VRAM", "LOW VRAM", "NORMAL VRAM", "HIGH VRAM", "MPS"][vram_state])
|
||||
print(f"Set vram state to: {vram_state.name}")
|
||||
|
||||
|
||||
current_loaded_model = None
|
||||
|
|
@ -109,12 +106,12 @@ def unload_model():
|
|||
model_accelerated = False
|
||||
|
||||
#never unload models from GPU on high vram
|
||||
if vram_state != HIGH_VRAM:
|
||||
if vram_state != VRAMState.HIGH_VRAM:
|
||||
current_loaded_model.model.cpu()
|
||||
current_loaded_model.unpatch_model()
|
||||
current_loaded_model = None
|
||||
|
||||
if vram_state != HIGH_VRAM:
|
||||
if vram_state != VRAMState.HIGH_VRAM:
|
||||
if len(current_gpu_controlnets) > 0:
|
||||
for n in current_gpu_controlnets:
|
||||
n.cpu()
|
||||
|
|
@ -135,19 +132,19 @@ def load_model_gpu(model):
|
|||
model.unpatch_model()
|
||||
raise e
|
||||
current_loaded_model = model
|
||||
if vram_state == CPU:
|
||||
if vram_state == VRAMState.CPU:
|
||||
pass
|
||||
elif vram_state == MPS:
|
||||
elif vram_state == VRAMState.MPS:
|
||||
mps_device = torch.device("mps")
|
||||
real_model.to(mps_device)
|
||||
pass
|
||||
elif vram_state == NORMAL_VRAM or vram_state == HIGH_VRAM:
|
||||
elif vram_state == VRAMState.NORMAL_VRAM or vram_state == VRAMState.HIGH_VRAM:
|
||||
model_accelerated = False
|
||||
real_model.cuda()
|
||||
else:
|
||||
if vram_state == NO_VRAM:
|
||||
if vram_state == VRAMState.NO_VRAM:
|
||||
device_map = accelerate.infer_auto_device_map(real_model, max_memory={0: "256MiB", "cpu": "16GiB"})
|
||||
elif vram_state == LOW_VRAM:
|
||||
elif vram_state == VRAMState.LOW_VRAM:
|
||||
device_map = accelerate.infer_auto_device_map(real_model, max_memory={0: "{}MiB".format(total_vram_available_mb), "cpu": "16GiB"})
|
||||
|
||||
accelerate.dispatch_model(real_model, device_map=device_map, main_device="cuda")
|
||||
|
|
@ -157,10 +154,10 @@ def load_model_gpu(model):
|
|||
def load_controlnet_gpu(models):
|
||||
global current_gpu_controlnets
|
||||
global vram_state
|
||||
if vram_state == CPU:
|
||||
if vram_state == VRAMState.CPU:
|
||||
return
|
||||
|
||||
if vram_state == LOW_VRAM or vram_state == NO_VRAM:
|
||||
if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
|
||||
#don't load controlnets like this if low vram because they will be loaded right before running and unloaded right after
|
||||
return
|
||||
|
||||
|
|
@ -176,20 +173,20 @@ def load_controlnet_gpu(models):
|
|||
|
||||
def load_if_low_vram(model):
|
||||
global vram_state
|
||||
if vram_state == LOW_VRAM or vram_state == NO_VRAM:
|
||||
if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
|
||||
return model.cuda()
|
||||
return model
|
||||
|
||||
def unload_if_low_vram(model):
|
||||
global vram_state
|
||||
if vram_state == LOW_VRAM or vram_state == NO_VRAM:
|
||||
if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
|
||||
return model.cpu()
|
||||
return model
|
||||
|
||||
def get_torch_device():
|
||||
if vram_state == MPS:
|
||||
if vram_state == VRAMState.MPS:
|
||||
return torch.device("mps")
|
||||
if vram_state == CPU:
|
||||
if vram_state == VRAMState.CPU:
|
||||
return torch.device("cpu")
|
||||
else:
|
||||
return torch.cuda.current_device()
|
||||
|
|
@ -201,9 +198,9 @@ def get_autocast_device(dev):
|
|||
|
||||
|
||||
def xformers_enabled():
|
||||
if vram_state == CPU:
|
||||
if vram_state == VRAMState.CPU:
|
||||
return False
|
||||
return XFORMERS_IS_AVAILBLE
|
||||
return XFORMERS_IS_AVAILABLE
|
||||
|
||||
|
||||
def xformers_enabled_vae():
|
||||
|
|
@ -243,7 +240,7 @@ def get_free_memory(dev=None, torch_free_too=False):
|
|||
|
||||
def maximum_batch_area():
|
||||
global vram_state
|
||||
if vram_state == NO_VRAM:
|
||||
if vram_state == VRAMState.NO_VRAM:
|
||||
return 0
|
||||
|
||||
memory_free = get_free_memory() / (1024 * 1024)
|
||||
|
|
@ -252,11 +249,11 @@ def maximum_batch_area():
|
|||
|
||||
def cpu_mode():
|
||||
global vram_state
|
||||
return vram_state == CPU
|
||||
return vram_state == VRAMState.CPU
|
||||
|
||||
def mps_mode():
|
||||
global vram_state
|
||||
return vram_state == MPS
|
||||
return vram_state == VRAMState.MPS
|
||||
|
||||
def should_use_fp16():
|
||||
if cpu_mode() or mps_mode():
|
||||
|
|
|
|||
27
main.py
27
main.py
|
|
@ -1,37 +1,14 @@
|
|||
import argparse
|
||||
import asyncio
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import threading
|
||||
from comfy.cli_args import args
|
||||
|
||||
if os.name == "nt":
|
||||
import logging
|
||||
logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage())
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Script Arguments")
|
||||
|
||||
parser.add_argument("--listen", type=str, default="127.0.0.1", help="Listen on IP or 0.0.0.0 if none given so the UI can be accessed from other computers.")
|
||||
parser.add_argument("--port", type=int, default=8188, help="Set the listen port.")
|
||||
parser.add_argument("--extra-model-paths-config", type=str, default=None, help="Load an extra_model_paths.yaml file.")
|
||||
parser.add_argument("--output-directory", type=str, default=None, help="Set the ComfyUI output directory.")
|
||||
parser.add_argument("--dont-upcast-attention", action="store_true", help="Disable upcasting of attention. Can boost speed but increase the chances of black images.")
|
||||
parser.add_argument("--use-split-cross-attention", action="store_true", help="Use the split cross attention optimization instead of the sub-quadratic one. Ignored when xformers is used.")
|
||||
parser.add_argument("--use-pytorch-cross-attention", action="store_true", help="Use the new pytorch 2.0 cross attention function.")
|
||||
parser.add_argument("--disable-xformers", action="store_true", help="Disable xformers.")
|
||||
parser.add_argument("--cuda-device", type=int, default=None, help="Set the id of the cuda device this instance will use.")
|
||||
parser.add_argument("--highvram", action="store_true", help="By default models will be unloaded to CPU memory after being used. This option keeps them in GPU memory.")
|
||||
parser.add_argument("--normalvram", action="store_true", help="Used to force normal vram use if lowvram gets automatically enabled.")
|
||||
parser.add_argument("--lowvram", action="store_true", help="Split the unet in parts to use less vram.")
|
||||
parser.add_argument("--novram", action="store_true", help="When lowvram isn't enough.")
|
||||
parser.add_argument("--cpu", action="store_true", help="To use the CPU for everything (slow).")
|
||||
parser.add_argument("--dont-print-server", action="store_true", help="Don't print server output.")
|
||||
parser.add_argument("--quick-test-for-ci", action="store_true", help="Quick test for CI.")
|
||||
parser.add_argument("--windows-standalone-build", action="store_true", help="Windows standalone build.")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.dont_upcast_attention:
|
||||
print("disabling upcasting of attention")
|
||||
os.environ['ATTN_PRECISION'] = "fp16"
|
||||
|
|
@ -121,7 +98,7 @@ if __name__ == "__main__":
|
|||
|
||||
if args.output_directory:
|
||||
output_dir = os.path.abspath(args.output_directory)
|
||||
print("setting output directory to:", output_dir)
|
||||
print(f"Setting output directory to: {output_dir}")
|
||||
folder_paths.set_output_directory(output_dir)
|
||||
|
||||
port = args.port
|
||||
|
|
|
|||
Loading…
Reference in New Issue