Fix oobabooga#5036

randoentity · Jul 27, 2024 · c76711c · c76711c
1 parent 6bab4c2
commit c76711c
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/modules/models.py b/modules/models.py
@@ -152,6 +152,15 @@ def huggingface_loader(model_name):
 
     if 'chatglm' in model_name.lower():
         LoaderClass = AutoModel
+    elif 'llava' in model_name.lower():
+        from transformers import LlamaConfig, LlamaForCausalLM
+        class LlavaConfig(LlamaConfig):
+            model_type = "llava"
+        class LlavaLlamaForCausalLM(LlamaForCausalLM):
+            config_class = LlavaConfig
+        AutoConfig.register("llava", LlavaConfig, exist_ok=True)
+        AutoModelForCausalLM.register(LlavaConfig, LlavaLlamaForCausalLM, exist_ok=True)
+        LoaderClass = LlavaLlamaForCausalLM
     else:
         if config.to_dict().get('is_encoder_decoder', False):
             LoaderClass = AutoModelForSeq2SeqLM