From 5b0bbfa6e8e5366d727148ecce461433d73582d5 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Mon, 30 Jan 2023 14:17:12 -0300
Subject: [PATCH] Clean up

---
 server.py | 29 ++++++++++-------------------
 1 file changed, 10 insertions(+), 19 deletions(-)

diff --git a/server.py b/server.py
index 4efac0c..29439fb 100644
--- a/server.py
+++ b/server.py
@@ -8,7 +8,6 @@ import json
 import io
 import base64
 import sys
-from sys import exit
 from pathlib import Path
 from PIL import Image
 import copy
@@ -69,8 +68,7 @@ settings = {
 }
 
 if args.settings is not None and Path(args.settings).exists():
-    with open(Path(args.settings), 'r') as f:
-        new_settings = json.load(f)
+    new_settings = json.loads(open(Path(args.settings), 'r').read())
     for item in new_settings:
         settings[item] = new_settings[item]
 
@@ -96,20 +94,14 @@ def load_model(model_name):
             settings.append("torch_dtype=torch.float32")
         else:
             settings.append("device_map='auto'")
-            if args.gpu_memory is not None:
-                if args.cpu_memory is not None:
-                    settings.append(f"max_memory={{0: '{args.gpu_memory}GiB', 'cpu': '{args.cpu_memory}GiB'}}")
-                else:
-                    settings.append(f"max_memory={{0: '{args.gpu_memory}GiB', 'cpu': '99GiB'}}")
+            settings.append("load_in_8bit=True" if args.load_in_8bit else "torch_dtype=torch.float16")
+
+            if args.gpu_memory and args.cpu_memory:
+                settings.append(f"max_memory={{0: '{args.gpu_memory}GiB', 'cpu': '{args.cpu_memory}GiB'}}")
+            elif args.gpu_memory:
+                settings.append(f"max_memory={{0: '{args.gpu_memory}GiB', 'cpu': '99GiB'}}")
             if args.disk:
-                if args.disk_cache_dir is not None:
-                    settings.append(f"offload_folder='{args.disk_cache_dir}'")
-                else:
-                    settings.append("offload_folder='cache'")
-            if args.load_in_8bit:
-                settings.append("load_in_8bit=True")
-            else:
-                settings.append("torch_dtype=torch.float16")
+                settings.append(f"offload_folder='{args.disk_cache_dir or 'cache'}'")
 
         settings = ', '.join(set(settings))
         command = f"{command}(Path(f'models/{model_name}'), {settings})"
@@ -307,7 +299,7 @@ if args.model is not None:
 else:
     if len(available_models) == 0:
         print("No models are available! Please download at least one.")
-        exit(0)
+        sys.exit(0)
     elif len(available_models) == 1:
         i = 0
     else:
@@ -566,8 +558,7 @@ if args.chat or args.cai_chat:
         history['visible'] = []
         if _character != 'None':
             character = _character
-            with open(Path(f'characters/{_character}.json'), 'r') as f:
-                data = json.loads(f.read())
+            data = json.loads(open(Path(f'characters/{_character}.json'), 'r').read())
             name2 = data['char_name']
             if 'char_persona' in data and data['char_persona'] != '':
                 context += f"{data['char_name']}'s Persona: {data['char_persona']}\n"