Update onnx/builder.py

onnx/builder.py

@@ -17,7 +17,7 @@ def build_vision(args):
     prompt = f"{user_prompt}<|image_1|>\nWhat is shown in this image?{prompt_suffix}{assistant_prompt}"
     url = "https://www.ilankelman.org/stopsigns/australia.jpg"
     image = Image.open(requests.get(url, stream=True).raw)
-    inputs = processor(prompt, image, return_tensors="pt").to(args.execution_provider)
+    inputs = processor(prompt, image, return_tensors="pt").to(args.execution_provider.replace("dml", "cuda"))
     inputs["pixel_values"] = inputs["pixel_values"].to(args.precision)
 
     # TorchScript export
@@ -214,8 +214,8 @@ def get_args():
         "-e",
         "--execution_provider",
         required=True,
-        choices=["cpu", "cuda"],
-        help="Device to export Phi-3 vision components with",
+        choices=["cpu", "cuda", "dml"],
+        help="Execution provider for Phi-3 vision components",
     )
 
     parser.add_argument(
@@ -238,7 +238,7 @@ if __name__ == "__main__":
     args = get_args()
     config = AutoConfig.from_pretrained(args.input, trust_remote_code=True)
     processor = AutoProcessor.from_pretrained(args.input, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(args.input, trust_remote_code=True, torch_dtype=args.precision).to(args.execution_provider)
+    model = AutoModelForCausalLM.from_pretrained(args.input, trust_remote_code=True, torch_dtype=args.precision).to(args.execution_provider.replace("dml", "cuda"))
 
     # Build model components
     build_vision(args)