Update handler.py

handler.py

@@ -1,7 +1,6 @@
 from typing import  Dict, List, Any
 from PIL import Image
 import torch
-from torch import autocast
 from diffusers import StableDiffusionUpscalePipeline
 import base64
 from io import BytesIO
@@ -24,9 +23,7 @@ class EndpointHandler():
         self.path = path
         # load the optimized model
         #model_id = "stabilityai/stable-diffusion-x4-upscaler"
-        self.pipe = StableDiffusionUpscalePipeline.from_pretrained(path, torch_dtype=torch.float16)
-        self.pipe = self.pipe.to(device)
-
+        
 
     def __call__(self, data) -> List[Dict[str, Any]]:
         """
@@ -35,6 +32,9 @@ class EndpointHandler():
         Return:
             A :obj:`dict`:. base64 encoded image
         """
+
+        self.pipe = StableDiffusionUpscalePipeline.from_pretrained(path, torch_dtype=torch.float16)
+        self.pipe = self.pipe.to(device)
         logger.info('data received %s', data)
         inputs = data.get("inputs")
         logger.info('inputs received %s', inputs)