Update app.py

adjust architecture for certain models

app.py

@@ -26,9 +26,33 @@ def get_model_summary(model_name):
         return model_cache[model_name], ""
 
     try:
+        # Fetch the config.json file
+        config_url = f"https://huggingface.co/{model_name}/raw/main/config.json"
+        response = requests.get(config_url)
+        response.raise_for_status()
+        config = response.json()
+        architecture = config["architectures"][0]
+
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
-        model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True).to(device)
+        # Select the correct model class based on the architecture
+        if architecture == "LlavaNextForConditionalGeneration":
+            from transformers import LlavaNextForConditionalGeneration
+            model = LlavaNextForConditionalGeneration.from_pretrained(model_name, trust_remote_code=True).to(device)
+        elif architecture == "LlavaForConditionalGeneration":
+            from transformers import LlavaForConditionalGeneration
+            model = LlavaForConditionalGeneration.from_pretrained(model_name, trust_remote_code=True).to(device)
+        elif architecture == "PaliGemmaForConditionalGeneration":
+            from transformers import PaliGemmaForConditionalGeneration
+            model = PaliGemmaForConditionalGeneration.from_pretrained(model_name, trust_remote_code=True).to(device)
+        elif architecture == "Idefics2ForConditionalGeneration":
+            from transformers import Idefics2ForConditionalGeneration
+            model = Idefics2ForConditionalGeneration.from_pretrained(model_name, trust_remote_code=True).to(device)
+        elif architecture == "MiniCPMV":
+            from transformers import MiniCPMV
+            model = MiniCPMV.from_pretrained(model_name, trust_remote_code=True).to(device)
+        else:
+            model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True).to(device)
         
         model_summary = str(model)
         model_cache[model_name] = model_summary