File size: 1,814 Bytes
2c83deb
d70b1c6
2c83deb
732c9a5
2c83deb
 
26fde58
 
 
 
 
 
2c83deb
 
 
 
 
7fd98b7
 
2c83deb
 
 
2c745bf
2c83deb
732c9a5
be1c648
2c83deb
be8f6aa
2c83deb
 
4fa2b44
2c83deb
 
 
be1c648
ecee85b
be1c648
43b0eb9
be8f6aa
 
 
2c83deb
01e3be6
 
 
 
7fd98b7
b4614a6
01e3be6
ff2ef5c
 
22d7300
ba24c0a
 
 
22d7300
2c83deb
d188461
ba24c0a
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
from typing import  Dict, List, Any
from PIL import Image
import torch
from diffusers import StableDiffusionUpscalePipeline
import base64
from io import BytesIO
from transformers.utils import logging

logging.set_verbosity_info()
logger = logging.get_logger("transformers")
logger.info("INFO")
logger.warning("WARN")


# set device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

#if device.type != 'cuda':
#    raise ValueError("need to run on GPU")

class EndpointHandler():
    def __init__(self, path=""):
        self.path = path
        # load the optimized model
        #model_id = "stabilityai/stable-diffusion-x4-upscaler"
        

    def __call__(self, data) -> List[Dict[str, Any]]:
        """
        Args:
            image (:obj:`string`)
        Return:
            A :obj:`dict`:. base64 encoded image
        """

        self.pipe = StableDiffusionUpscalePipeline.from_pretrained(self.path, torch_dtype=torch.float16)
        self.pipe = self.pipe.to(device)
        #self.pipe.enable_xformers_memory_efficient_attention()
        logger.info('data received %s', data)
        inputs = data.get("inputs")
        logger.info('inputs received %s', inputs)

        image_base64 = base64.b64decode(inputs['image'])
        logger.info('image_base64')
        image_bytes = BytesIO(image_base64)
        logger.info('image_bytes')
        image = Image.open(image_bytes).convert("RGB")
        prompt = inputs['prompt']
        logger.info('image')
        
        upscaled_image = self.pipe(prompt, image).images[0]

        buffered = BytesIO()
        upscaled_image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue())

        # postprocess the prediction
        return {"image": img_str.decode()}
        #return {"image": "test"}