Update app.py

app.py

@@ -199,13 +199,9 @@ def run_inference(prompt, image, temperature=0.2, top_p=0.95, max_tokens=512):
             text_prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
             inputs = processor(text=text_prompt, images=image, return_tensors="pt")
 
-        target_dtype = getattr(model, "dtype", None)
         processed_inputs = {}
         for key, value in inputs.items():
-            if torch.is_floating_point(value) and target_dtype is not None:
-                processed_inputs[key] = value.to(model.device, dtype=target_dtype)
-            else:
-                processed_inputs[key] = value.to(model.device)
+            processed_inputs[key] = value.to(model.device, dtype=torch.bfloat16)
         inputs = processed_inputs
 
         with torch.no_grad():
@@ -596,7 +592,7 @@ logger.info(f"Loading model: {model_path}")
 if "qwenguard" in model_path.lower():
     model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
         model_path,
-        torch_dtype="auto",
+        torch_dtype=torch.bfloat16,
         device_map="auto"
     )
     processor = AutoProcessor.from_pretrained(model_path)
@@ -606,7 +602,7 @@ if "qwenguard" in model_path.lower():
 else:
     model = LlavaOnevisionForConditionalGeneration.from_pretrained(
         model_path,
-        torch_dtype="auto",
+        torch_dtype=torch.bfloat16,
         device_map="auto",
         trust_remote_code=True
     )