Commit
•
41339c4
1
Parent(s):
dcc0b95
Update handler.py
Browse files- handler.py +2 -1
handler.py
CHANGED
|
@@ -6,6 +6,7 @@ import base64
|
|
| 6 |
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
|
| 7 |
import logging
|
| 8 |
import subprocess
|
|
|
|
| 9 |
|
| 10 |
logger = logging.getLogger()
|
| 11 |
logger.setLevel(logging.DEBUG)
|
|
@@ -23,7 +24,6 @@ class EndpointHandler():
|
|
| 23 |
self.pipe.to(device)
|
| 24 |
|
| 25 |
logger.info(f"PIPE LOADED")
|
| 26 |
-
subprocess.run("nvidia-smi")
|
| 27 |
|
| 28 |
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
|
| 29 |
"""
|
|
@@ -44,6 +44,7 @@ class EndpointHandler():
|
|
| 44 |
|
| 45 |
parameters = data.pop("parameters", data)
|
| 46 |
prompt = parameters['prompt']
|
|
|
|
| 47 |
|
| 48 |
images = self.pipe(prompt, image=image, num_inference_steps=10, image_guidance_scale=1).images
|
| 49 |
|
|
|
|
| 6 |
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
|
| 7 |
import logging
|
| 8 |
import subprocess
|
| 9 |
+
import sys
|
| 10 |
|
| 11 |
logger = logging.getLogger()
|
| 12 |
logger.setLevel(logging.DEBUG)
|
|
|
|
| 24 |
self.pipe.to(device)
|
| 25 |
|
| 26 |
logger.info(f"PIPE LOADED")
|
|
|
|
| 27 |
|
| 28 |
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
|
| 29 |
"""
|
|
|
|
| 44 |
|
| 45 |
parameters = data.pop("parameters", data)
|
| 46 |
prompt = parameters['prompt']
|
| 47 |
+
subprocess.run("nvidia-smi")
|
| 48 |
|
| 49 |
images = self.pipe(prompt, image=image, num_inference_steps=10, image_guidance_scale=1).images
|
| 50 |
|