bfloat16

app.py

@@ -101,13 +101,23 @@ def get_som_response(instruction, image_som):
         add_generation_prompt=True
     )
 
+    # inputs = magma_processor(images=[image_som], texts=prompt, return_tensors="pt")
+    # # with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+    # #     inputs['pixel_values'] = inputs['pixel_values'].unsqueeze(0).to(torch.bfloat16) # Add .to(torch.bfloat16) here for explicit casting
+    # #     inputs['image_sizes'] = inputs['image_sizes'].unsqueeze(0)
+    # #     logger.warning(inputs['pixel_values'].dtype)
+    # #     # inputs = inputs.to("cuda")
+    # inputs = inputs.to("cuda", dtype=torch.bfloat16)
+
     inputs = magma_processor(images=[image_som], texts=prompt, return_tensors="pt")
-    # with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
-    #     inputs['pixel_values'] = inputs['pixel_values'].unsqueeze(0).to(torch.bfloat16) # Add .to(torch.bfloat16) here for explicit casting
-    #     inputs['image_sizes'] = inputs['image_sizes'].unsqueeze(0)
-    #     logger.warning(inputs['pixel_values'].dtype)
-    #     # inputs = inputs.to("cuda")
-    inputs = inputs.to("cuda", dtype=torch.bfloat16)
+    inputs['pixel_values'] = inputs['pixel_values'].to("cuda", dtype=torch.bfloat16)
+    inputs['image_sizes'] = inputs['image_sizes'].unsqueeze(0)
+    inputs['image_sizes'] = inputs['image_sizes'].to("cuda")
+
+    # 处理其他可能的输入
+    for key in inputs:
+        if key not in ['pixel_values', 'image_sizes'] and torch.is_tensor(inputs[key]):
+            inputs[key] = inputs[key].to("cuda")
 
     magam_model.generation_config.pad_token_id = magma_processor.tokenizer.pad_token_id
     with torch.inference_mode():