From a1e4405aa68635cd5abcd1e56c5dfb7f2a5cc16c Mon Sep 17 00:00:00 2001
From: Gnome Ann <>
Date: Tue, 5 Oct 2021 10:36:51 -0400
Subject: [PATCH] Automatically use breakmodel instead of GPU-only where
 supported

There's really no reason to use GPU-only mode if breakmodel is supported
because breakmodel can run in GPU-only mode too.
---
 aiserver.py | 26 +++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/aiserver.py b/aiserver.py
index ddce311d..c3dd4292 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -314,12 +314,12 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
             genselected = True
             vars.usegpu = True
             vars.breakmodel = False
+        if(vars.bmsupported):
+            vars.usegpu = False
+            vars.breakmodel = True
         if(args.cpu):
             vars.usegpu = False
             vars.breakmodel = False
-        if(vars.bmsupported and args.breakmodel):
-            vars.usegpu = False
-            vars.breakmodel = True
     elif(vars.hascuda):    
         if(vars.bmsupported):
             print(colors.YELLOW + "You're using a model that supports hybrid generation!")
@@ -327,9 +327,8 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
             print("(slower than GPU-only but uses less VRAM) or between multiple GPUs")
             print("(allowing you to use the combined VRAM of all your GPUs).")
             print("Currently only GPT-Neo and GPT-J models support this feature.")
-        print("{0}Use GPU or CPU for generation?:  (Default GPU){1}".format(colors.CYAN, colors.END))
-        if(vars.bmsupported):
-            print(f"    1 - GPU\n    2 - CPU\n    3 - Hybrid generation\n")
+            print("{0}Use hybrid generation or CPU-only generation?:  (Default hybrid){1}".format(colors.CYAN, colors.END))
+            print(f"    1 - Hybrid generation\n    2 - CPU\n")
         else:
             print("    1 - GPU\n    2 - CPU\n")
         genselected = False
@@ -342,17 +341,18 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
                 vars.usegpu = True
                 genselected = True
             elif(genselect.isnumeric() and int(genselect) == 1):
-                vars.breakmodel = False
-                vars.usegpu = True
-                genselected = True
+                if(vars.bmsupported):
+                    vars.breakmodel = True
+                    vars.usegpu = False
+                    genselected = True
+                else:
+                    vars.breakmodel = False
+                    vars.usegpu = True
+                    genselected = True
             elif(genselect.isnumeric() and int(genselect) == 2):
                 vars.breakmodel = False
                 vars.usegpu = False
                 genselected = True
-            elif(vars.bmsupported and genselect.isnumeric() and int(genselect) == 3):
-                vars.breakmodel = True
-                vars.usegpu = False
-                genselected = True
             else:
                 print("{0}Please enter a valid selection.{1}".format(colors.RED, colors.END))