Fix tool loading errors in MedRAX2

- Fix NV-Reason-CXR quantization config handling to avoid dict.to_dict() error
- Correct tool import paths to use actual class names:
- CheXagentXRayVQATool instead of XRayVQATool
- TorchXRayVisionClassifierTool instead of XRayClassificationTool
- ChestXRayReportGeneratorTool instead of XRayReportGenerationTool
- DicomProcessorTool instead of DICOMTool
- WebBrowserTool instead of WebBrowsingTool
- Add better error handling for model loading

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

app.py

@@ -58,8 +58,8 @@ if device == "cuda":
         print(f"✗ Failed to load grounding tool: {e}")
 
     try:
-        from medrax.tools import XRayVQATool
-        vqa_tool = XRayVQATool(
+        from medrax.tools.vqa import CheXagentXRayVQATool
+        vqa_tool = CheXagentXRayVQATool(
             device=device,
             temp_dir="temp",
             load_in_4bit=True
@@ -70,8 +70,8 @@ if device == "cuda":
         print(f"✗ Failed to load VQA tool: {e}")
 
     try:
-        from medrax.tools import XRayClassificationTool
-        classification_tool = XRayClassificationTool(
+        from medrax.tools.classification import TorchXRayVisionClassifierTool
+        classification_tool = TorchXRayVisionClassifierTool(
             device=device,
             temp_dir="temp",
             load_in_4bit=True
@@ -82,8 +82,8 @@ if device == "cuda":
         print(f"✗ Failed to load classification tool: {e}")
 
     try:
-        from medrax.tools import XRayReportGenerationTool
-        report_tool = XRayReportGenerationTool(
+        from medrax.tools.report_generation import ChestXRayReportGeneratorTool
+        report_tool = ChestXRayReportGeneratorTool(
             device=device,
             temp_dir="temp",
             load_in_4bit=True
@@ -95,16 +95,16 @@ if device == "cuda":
 
 # Load non-GPU tools
 try:
-    from medrax.tools import DICOMTool
-    dicom_tool = DICOMTool(temp_dir="temp")
+    from medrax.tools.dicom import DicomProcessorTool
+    dicom_tool = DicomProcessorTool(temp_dir="temp")
     tools.append(dicom_tool)
     print("✓ Loaded DICOM tool")
 except Exception as e:
     print(f"✗ Failed to load DICOM tool: {e}")
 
 try:
-    from medrax.tools import WebBrowsingTool
-    browsing_tool = WebBrowsingTool()
+    from medrax.tools.browsing import WebBrowserTool
+    browsing_tool = WebBrowserTool()
     tools.append(browsing_tool)
     print("✓ Loaded web browsing tool")
 except Exception as e:

medrax/tools/nv_reason_cxr.py

@@ -67,35 +67,50 @@ class NVReasonCXRTool(BaseTool):
         self.device = device
 
         # Setup quantization config
-        if load_in_4bit:
-            quantization_config = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_compute_dtype=torch.bfloat16,
-                bnb_4bit_use_double_quant=True,
-                bnb_4bit_quant_type="nf4",
-            )
-        else:
-            quantization_config = None
+        quantization_config = None
+        if load_in_4bit and device == "cuda":
+            try:
+                quantization_config = BitsAndBytesConfig(
+                    load_in_4bit=True,
+                    bnb_4bit_compute_dtype=torch.bfloat16,
+                    bnb_4bit_use_double_quant=True,
+                    bnb_4bit_quant_type="nf4",
+                )
+            except Exception as e:
+                print(f"Warning: Could not setup 4-bit quantization: {e}")
+                quantization_config = None
 
         # Load model
-        print(f"Loading NV-Reason-CXR model from {model_path}...")
-        self.model = AutoModelForImageTextToText.from_pretrained(
-            model_path,
-            device_map=self.device,
-            cache_dir=cache_dir,
-            torch_dtype=torch.bfloat16,
-            quantization_config=quantization_config,
-            trust_remote_code=True,
-        ).eval()
-
-        self.processor = AutoProcessor.from_pretrained(
-            model_path,
-            cache_dir=cache_dir,
-            trust_remote_code=True,
-            use_fast=True,
-        )
-
-        print(f"✓ NV-Reason-CXR model loaded successfully")
+        try:
+            print(f"Loading NV-Reason-CXR model from {model_path}...")
+
+            # Load without quantization config if it's causing issues
+            model_kwargs = {
+                "device_map": self.device,
+                "cache_dir": cache_dir,
+                "torch_dtype": torch.bfloat16,
+                "trust_remote_code": True,
+            }
+
+            if quantization_config is not None:
+                model_kwargs["quantization_config"] = quantization_config
+
+            self.model = AutoModelForImageTextToText.from_pretrained(
+                model_path,
+                **model_kwargs
+            ).eval()
+
+            self.processor = AutoProcessor.from_pretrained(
+                model_path,
+                cache_dir=cache_dir,
+                trust_remote_code=True,
+                use_fast=True,
+            )
+
+            print(f"✓ NV-Reason-CXR model loaded successfully")
+        except Exception as e:
+            print(f"Error loading NV-Reason-CXR model: {e}")
+            raise
 
     def _run(
         self,