Update handler.py

handler.py

@@ -6,6 +6,7 @@ import base64
 from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
 import logging
 import subprocess
+import sys
 
 logger = logging.getLogger()
 logger.setLevel(logging.DEBUG)
@@ -23,7 +24,6 @@ class EndpointHandler():
         self.pipe.to(device)
         
         logger.info(f"PIPE LOADED")
-        subprocess.run("nvidia-smi")
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -44,6 +44,7 @@ class EndpointHandler():
 
         parameters = data.pop("parameters", data)
         prompt = parameters['prompt']
+        subprocess.run("nvidia-smi")
 
         images = self.pipe(prompt, image=image, num_inference_steps=10, image_guidance_scale=1).images