Merge pull request #402 from NexaAI/perry/server-dev

fixed flux nexa list problem
NexaAI · Feb 24, 2025 · 2c759da · 2c759da
2 parents 4b2997c + 9aac34b
commit 2c759da
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 4 deletions.
diff --git a/nexa/constants.py b/nexa/constants.py
@@ -3,8 +3,7 @@
 from enum import Enum
 
 # Paths for caching, model hub, and tokens
-NEXA_CACHE_ROOT = Path(os.getenv("NEXA_CACHE_ROOT")
-                       or "~/.cache/nexa").expanduser()
+NEXA_CACHE_ROOT = Path(os.getenv("NEXA_CACHE_ROOT") or "~/.cache/nexa").expanduser()
 NEXA_TOKEN_PATH = NEXA_CACHE_ROOT / "token"
 NEXA_MODELS_HUB_DIR = NEXA_CACHE_ROOT / "hub"
 NEXA_MODEL_EVAL_RESULTS_PATH = NEXA_CACHE_ROOT / "eval"
@@ -266,10 +265,15 @@ class ModelType(Enum):
 NEXA_RUN_MODEL_MAP_FLUX = {
     "flux": "FLUX.1-schnell:flux1-schnell-q4_0",
     "FLUX.1-schnell:q4_0": "FLUX.1-schnell:flux1-schnell-q4_0",
+    "FLUX.1-schnell:flux1-schnell-q4_0": "FLUX.1-schnell:flux1-schnell-q4_0",
     "FLUX.1-schnell:q5_0": "FLUX.1-schnell:flux1-schnell-q5_0",
+    "FLUX.1-schnell:flux1-schnell-q5_0": "FLUX.1-schnell:flux1-schnell-q5_0",
     "FLUX.1-schnell:q5_1": "FLUX.1-schnell:flux1-schnell-q5_1",
+    "FLUX.1-schnell:flux1-schnell-q5_1": "FLUX.1-schnell:flux1-schnell-q5_1",
     "FLUX.1-schnell:q8_0": "FLUX.1-schnell:flux1-schnell-q8_0",
+    "FLUX.1-schnell:flux1-schnell-q8_0": "FLUX.1-schnell:flux1-schnell-q8_0",
     "FLUX.1-schnell:fp16": "FLUX.1-schnell:flux1-schnell-fp16",
+    "FLUX.1-schnell:flux1-schnell-fp16": "FLUX.1-schnell:flux1-schnell-fp16",
 }
 
 NEXA_RUN_MODEL_MAP_TEXT_EMBEDDING = {
@@ -359,8 +363,7 @@ class ModelType(Enum):
 EXIT_COMMANDS = ["/exit", "/quit", "/bye"]
 EXIT_REMINDER = f"Please use Ctrl + d or one of {EXIT_COMMANDS} to exit.\n"
 
-NEXA_STOP_WORDS_MAP = {
-    "octopus-v2": ["<nexa_end>"], "octopus-v4": ["<nexa_end>"]}
+NEXA_STOP_WORDS_MAP = {"octopus-v2": ["<nexa_end>"], "octopus-v4": ["<nexa_end>"]}
 
 DEFAULT_TEXT_GEN_PARAMS = {
     "temperature": 0.7,

diff --git a/nexa/general.py b/nexa/general.py
@@ -679,6 +679,17 @@ def add_model_to_list(model_name, model_location, model_type, run_type):
         if tag_name.startswith("model-"):
             tag_name = tag_name[6:]
             model_name = f"{model_name.split(':')[0]}:{tag_name}"
+        else:
+            return
+
+    # For Computer Vision Flux model, should remove the "flux1-schnell-" prefix from the tag name
+    if run_type == "Computer Vision":
+        tag_name = model_name.split(":")[1]
+        if tag_name.startswith("flux1-schnell-"):
+            tag_name = tag_name[14:]
+            model_name = f"{model_name.split(':')[0]}:{tag_name}"
+        else:
+            return
 
     model_list[model_name] = {
         "type": model_type,