mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Update.

Browse source
This commit is contained in:
Silver267 2023-01-20 17:05:43 -05:00
parent c0f2367b54
commit f4634e4c32
2 changed files with 12 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

12
server.py
View file

@ -31,6 +31,8 @@ parser.add_argument('--no-stream', action='store_true', help='Don\'t stream the
parser.add_argument('--settings', type=str, help='Load the default interface settings from this json file. See settings-template.json for an example.')
parser.add_argument('--no-listen', action='store_true', help='Make the web UI unreachable from your local network.')
parser.add_argument('--share', action='store_true', help='Create a public URL. This is useful for running the web UI on Google Colab or similar.')
parser.add_argument('--max-cpu-mem', type=int, help='Maximum cpu memory in GiB to allocate to the memory for offloading.')
parser.add_argument('--disk-cache-dir', type=str, help='Directory which you want the disk cache to load to.')
args = parser.parse_args()
loaded_preset = None
@ -90,9 +92,15 @@ def load_model(model_name):
else:
settings.append("device_map='auto'")
if args.max_gpu_memory is not None:
settings.append(f"max_memory={{0: '{args.max_gpu_memory}GiB', 'cpu': '99GiB'}}")
if args.max_cpu_mem is not None:
settings.append(f"max_memory={{0: '{args.max_gpu_memory}GiB', 'cpu': '{args.max_cpu_mem}GiB'}}")
else:
settings.append(f"max_memory={{0: '{args.max_gpu_memory}GiB', 'cpu': '99GiB'}}")
if args.disk:
settings.append("offload_folder='cache'")
if args.disk_cache_dir is not None:
settings.append("offload_folder='"+args.disk_cache_dir+"'")
else:
settings.append("offload_folder='cache'")
if args.load_in_8bit:
settings.append("load_in_8bit=True")
else: