samwell Claude commited on
Commit
4df7c2f
·
1 Parent(s): e3751ec

Fix multiple tool loading issues

Browse files

- Remove all quantization code from NV-Reason-CXR to avoid dict.to_dict() error
- Fix Classification and Report Generation tools by removing unsupported temp_dir parameter
- Add tensorflow>=2.15.0 to requirements for VQA tool
- NV-Reason-CXR will now run at bfloat16 precision (~6GB memory)

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

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

Files changed (3) hide show
  1. app.py +0 -2
  2. medrax/tools/nv_reason_cxr.py +7 -27
  3. requirements.txt +1 -0
app.py CHANGED
@@ -73,7 +73,6 @@ if device == "cuda":
73
  from medrax.tools.classification import TorchXRayVisionClassifierTool
74
  classification_tool = TorchXRayVisionClassifierTool(
75
  device=device,
76
- temp_dir="temp",
77
  load_in_4bit=True
78
  )
79
  tools.append(classification_tool)
@@ -85,7 +84,6 @@ if device == "cuda":
85
  from medrax.tools.report_generation import ChestXRayReportGeneratorTool
86
  report_tool = ChestXRayReportGeneratorTool(
87
  device=device,
88
- temp_dir="temp",
89
  load_in_4bit=True
90
  )
91
  tools.append(report_tool)
 
73
  from medrax.tools.classification import TorchXRayVisionClassifierTool
74
  classification_tool = TorchXRayVisionClassifierTool(
75
  device=device,
 
76
  load_in_4bit=True
77
  )
78
  tools.append(classification_tool)
 
84
  from medrax.tools.report_generation import ChestXRayReportGeneratorTool
85
  report_tool = ChestXRayReportGeneratorTool(
86
  device=device,
 
87
  load_in_4bit=True
88
  )
89
  tools.append(report_tool)
medrax/tools/nv_reason_cxr.py CHANGED
@@ -66,38 +66,18 @@ class NVReasonCXRTool(BaseTool):
66
  super().__init__()
67
  self.device = device
68
 
69
- # Setup quantization config
70
- quantization_config = None
71
- if load_in_4bit and device == "cuda":
72
- try:
73
- quantization_config = BitsAndBytesConfig(
74
- load_in_4bit=True,
75
- bnb_4bit_compute_dtype=torch.bfloat16,
76
- bnb_4bit_use_double_quant=True,
77
- bnb_4bit_quant_type="nf4",
78
- )
79
- except Exception as e:
80
- print(f"Warning: Could not setup 4-bit quantization: {e}")
81
- quantization_config = None
82
-
83
- # Load model
84
  try:
85
  print(f"Loading NV-Reason-CXR model from {model_path}...")
86
 
87
- # Load without quantization config if it's causing issues
88
- model_kwargs = {
89
- "device_map": self.device,
90
- "cache_dir": cache_dir,
91
- "torch_dtype": torch.bfloat16,
92
- "trust_remote_code": True,
93
- }
94
-
95
- if quantization_config is not None:
96
- model_kwargs["quantization_config"] = quantization_config
97
-
98
  self.model = AutoModelForImageTextToText.from_pretrained(
99
  model_path,
100
- **model_kwargs
 
 
 
101
  ).eval()
102
 
103
  self.processor = AutoProcessor.from_pretrained(
 
66
  super().__init__()
67
  self.device = device
68
 
69
+ # Load model without quantization to avoid compatibility issues
 
 
 
 
 
 
 
 
 
 
 
 
 
 
70
  try:
71
  print(f"Loading NV-Reason-CXR model from {model_path}...")
72
 
73
+ # Note: Skipping quantization due to dict.to_dict() error
74
+ # Model will run at bfloat16 precision (~6GB memory)
 
 
 
 
 
 
 
 
 
75
  self.model = AutoModelForImageTextToText.from_pretrained(
76
  model_path,
77
+ device_map=self.device,
78
+ cache_dir=cache_dir,
79
+ torch_dtype=torch.bfloat16,
80
+ trust_remote_code=True,
81
  ).eval()
82
 
83
  self.processor = AutoProcessor.from_pretrained(
requirements.txt CHANGED
@@ -47,6 +47,7 @@ openai>=0.27.0
47
  backoff>=1.10.0
48
  torch>=2.2.0
49
  torchvision>=0.10.0
 
50
  scikit-image>=0.18.0
51
  opencv-python>=4.8.0
52
  matplotlib>=3.8.0
 
47
  backoff>=1.10.0
48
  torch>=2.2.0
49
  torchvision>=0.10.0
50
+ tensorflow>=2.15.0
51
  scikit-image>=0.18.0
52
  opencv-python>=4.8.0
53
  matplotlib>=3.8.0