oobabooga 2 лет назад
Родитель
Сommit
bcea196c9d
2 измененных файлов с 5 добавлено и 2 удалено
  1. 1 1
      requirements.txt
  2. 4 1
      server.py

+ 1 - 1
requirements.txt

@@ -1,6 +1,6 @@
 accelerate==0.16.0
 accelerate==0.16.0
 bitsandbytes==0.37.0
 bitsandbytes==0.37.0
-flexgen==0.1.6
+flexgen==0.1.7
 gradio==3.18.0
 gradio==3.18.0
 numpy
 numpy
 rwkv==0.0.6
 rwkv==0.0.6

+ 4 - 1
server.py

@@ -28,7 +28,10 @@ if shared.args.settings is not None and Path(shared.args.settings).exists():
         shared.settings[item] = new_settings[item]
         shared.settings[item] = new_settings[item]
 
 
 def get_available_models():
 def get_available_models():
-    return sorted([item.name for item in list(Path('models/').glob('*')) if not item.name.endswith(('.txt', '-np'))], key=str.lower)
+    if shared.args.flexgen:
+        return sorted([re.sub('-np$', '', item.name) for item in list(Path('models/').glob('*')) if item.name.endswith('-np')], key=str.lower)
+    else:
+        return sorted([item.name for item in list(Path('models/').glob('*')) if not item.name.endswith(('.txt', '-np'))], key=str.lower)
 
 
 def get_available_presets():
 def get_available_presets():
     return sorted(set(map(lambda x : '.'.join(str(x.name).split('.')[:-1]), Path('presets').glob('*.txt'))), key=str.lower)
     return sorted(set(map(lambda x : '.'.join(str(x.name).split('.')[:-1]), Path('presets').glob('*.txt'))), key=str.lower)