Update app.py

app.py

@@ -11,10 +11,10 @@ import gradio as gr
 import torch
 from PIL import Image
 from transformers import (
-    AutoModelForCausalLM, 
-    AutoProcessor, 
+    AutoProcessor,
     AutoTokenizer, 
-    Qwen2_5_VLForConditionalGeneration
+    Qwen2_5_VLForConditionalGeneration,
+    LlavaOnevisionForConditionalGeneration
 )
 
 from taxonomy import policy_v1
@@ -113,7 +113,7 @@ default_conversation = Conversation()
 tokenizer = None
 model = None
 processor = None
-context_len = 2048
+context_len = 8048
 
 # Helper functions
 def clear_conv(conv):
@@ -150,7 +150,7 @@ def load_model(model_path):
             
         # Otherwise assume it's a LlavaGuard model
         else:
-            model = AutoModelForCausalLM.from_pretrained(
+            model = LlavaOnevisionForConditionalGeneration.from_pretrained(
                 model_path,
                 torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                 device_map="auto" if torch.cuda.is_available() else None,
@@ -159,7 +159,7 @@ def load_model(model_path):
             tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
             processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
         
-        context_len = getattr(model.config, "max_position_embeddings", 2048)
+        context_len = getattr(model.config, "max_position_embeddings", 8048)
         logger.info(f"Model {model_path} loaded successfully")
         return True
     
@@ -169,10 +169,10 @@ def load_model(model_path):
 
 def get_model_list():
     models = [
-        'AIML-TUDA/LlavaGuard-v1.2-7B-OV-hf',
         'AIML-TUDA/LlavaGuard-v1.2-0.5B-OV-hf',
+        'AIML-TUDA/LlavaGuard-v1.2-7B-OV-hf',
+        'AIML-TUDA/QwenGuard-v1.2-3B',
         'AIML-TUDA/QwenGuard-v1.2-7B',
-        'AIML-TUDA/QwenGuard-v1.2-3B'
     ]
     return models
 
@@ -238,12 +238,22 @@ def run_inference(prompt, image, temperature=0.2, top_p=0.95, max_tokens=512):
             
         # Otherwise assume it's a LlavaGuard model
         else:
+            conversation = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image"},
+                        {"type": "text", "text": prompt},
+                    ],
+                },
+            ]
+
+            text_prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+
+
             # Process input for LlavaGuard models
-            inputs = processor(
-                prompt, 
-                images=image, 
-                return_tensors="pt"
-            )
+            inputs = processor(text=text_prompt, images=image, return_tensors="pt")
+
             
             # Move to GPU if available
             if torch.cuda.is_available():