Re-enable GPU-only generation option

This commit is contained in:
Gnome Ann 2021-11-14 18:24:51 -05:00
parent 3e38b462c6
commit b0ab30cec4
1 changed files with 6 additions and 2 deletions

View File

@ -359,8 +359,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 hybrid generation or CPU-only generation?: (Default hybrid){1}".format(colors.CYAN, colors.END))
print(f" 1 - Hybrid generation\n 2 - CPU\n")
print("{0}Use hybrid generation, GPU-only generation or CPU-only generation?: (Default hybrid){1}".format(colors.CYAN, colors.END))
print(f" 1 - Hybrid generation\n 2 - GPU\n 3 - CPU\n")
else:
print(" 1 - GPU\n 2 - CPU\n")
genselected = False
@ -382,6 +382,10 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
vars.usegpu = True
genselected = True
elif(genselect.isnumeric() and int(genselect) == 2):
vars.breakmodel = False
vars.usegpu = True
genselected = True
elif(genselect.isnumeric() and int(genselect) == 3):
vars.breakmodel = False
vars.usegpu = False
genselected = True