Update handler.py

handler.py

@@ -6,8 +6,8 @@ from transformers import AutoProcessor, AutoModelForCausalLM
 
 class EndpointHandler():
     def __init__(self, path=""):
-        device = "cuda:0" if torch.cuda.is_available() else "cpu"
-        torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+        self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        self.torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
         self.model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-large", torch_dtype=torch_dtype, trust_remote_code=True).to(device)
         self.processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True)
@@ -17,7 +17,7 @@ class EndpointHandler():
         url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true"
         image = Image.open(requests.get(url, stream=True).raw)
 
-        inputs = self.processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
+        inputs = self.processor(text=prompt, images=image, return_tensors="pt").to(self.device, self.torch_dtype)
 
         generated_ids = self.model.generate(
             input_ids=inputs["input_ids"],