update

eagle_vl/serve/inference.py

@@ -18,7 +18,7 @@ from .chat_utils import Conversation, get_conv_template
 logger = logging.getLogger(__name__)
 
 
-def load_model(model_path: str = "nvidia/Eagle-2.5-8B"):
+def load_model_from_nv(model_path: str = "nvidia/Eagle-2.5-8B"):
 
     token = os.environ.get("HF_TOKEN")
     # hotfix the model to use flash attention 2
@@ -41,6 +41,36 @@ def load_model(model_path: str = "nvidia/Eagle-2.5-8B"):
 
     return model, processor
 
+def load_model_from_eagle(model_path: str = "NVEagle/Eagle2.5-VL-8B-Preview"):
+
+    token = os.environ.get("HF_TOKEN")
+    # hotfix the model to use flash attention 2
+    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True, use_auth_token=token)
+    config._attn_implementation = "flash_attention_2"
+    config.vision_config._attn_implementation = "flash_attention_2"
+    config.text_config._attn_implementation = "flash_attention_2"
+    print("Successfully set the attn_implementation to flash_attention_2")
+
+    logger.info(f"token = {token[:4]}***{token[-2:]}")
+    model = AutoModel.from_pretrained(
+        model_path,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+        use_auth_token=token
+    )
+    model.to("cuda")
+    processor = AutoProcessor.from_pretrained(model_path, config=config, trust_remote_code=True, use_fast=True, use_auth_token=token)
+
+    return model, processor
+
+def load_model(model_path: str = "nvidia/Eagle-2.5-8B"):
+    try:
+        model, processor = load_model_from_nv(model_path)
+    except Exception as e:
+        logger.error(f"Failed to load model from HF, trying to load from eagle: {e}")
+        model, processor = load_model_from_eagle(model_path)
+    return model, processor
 
 class StoppingCriteriaSub(StoppingCriteria):
     def __init__(self, stops=[], encounters=1):