HieuPM commited on
Commit
c6aeea5
·
1 Parent(s): 493a2dd

EDIT: remove super-resolution model and add inpainting sd

Browse files
Files changed (2) hide show
  1. handler.py +10 -10
  2. requirements.txt +0 -1
handler.py CHANGED
@@ -5,7 +5,7 @@ from PIL import Image
5
  import base64
6
  from io import BytesIO
7
  import numpy as np
8
- from RealESRGAN import RealESRGAN
9
 
10
  # set device
11
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -17,15 +17,15 @@ class EndpointHandler():
17
  def __init__(self, path=""):
18
  # load StableDiffusionInpaintPipeline pipeline
19
  self.pipe = StableDiffusionInpaintPipeline.from_pretrained(path, torch_dtype=torch.float16)
20
- # use DPMSolverMultistepScheduler
21
- self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
22
  # pipe.enable_sequential_cpu_offload()
23
  # move to device
24
  self.pipe.to(device)
25
  self.pipe.enable_xformers_memory_efficient_attention()
26
 
27
- self.upscaler = RealESRGAN(device, scale=4)
28
- self.upscaler.load_weights('weights/RealESRGAN_x4.pth', download=True)
29
 
30
 
31
  def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
@@ -68,11 +68,11 @@ class EndpointHandler():
68
  width=width
69
  ).images
70
 
71
- for i in range(len(out)):
72
- gen_img = Image.composite(out[i], image.resize(out[i].size), mask_image.resize(out[i].size))
73
- gen_img = self.upscaler.predict(gen_img)
74
- gen_img = Image.composite(gen_img, image.resize(gen_img.size), mask_image.resize(gen_img.size))
75
- out[i] = gen_img
76
 
77
  # return first generate PIL image
78
  json_imgs = {}
 
5
  import base64
6
  from io import BytesIO
7
  import numpy as np
8
+ # from RealESRGAN import RealESRGAN
9
 
10
  # set device
11
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
17
  def __init__(self, path=""):
18
  # load StableDiffusionInpaintPipeline pipeline
19
  self.pipe = StableDiffusionInpaintPipeline.from_pretrained(path, torch_dtype=torch.float16)
20
+ # use EulerAncestralDiscreteScheduler
21
+ self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config)
22
  # pipe.enable_sequential_cpu_offload()
23
  # move to device
24
  self.pipe.to(device)
25
  self.pipe.enable_xformers_memory_efficient_attention()
26
 
27
+ # self.upscaler = RealESRGAN(device, scale=4)
28
+ # self.upscaler.load_weights('weights/RealESRGAN_x4.pth', download=True)
29
 
30
 
31
  def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
 
68
  width=width
69
  ).images
70
 
71
+ # for i in range(len(out)):
72
+ # gen_img = Image.composite(out[i], image.resize(out[i].size), mask_image.resize(out[i].size))
73
+ # gen_img = self.upscaler.predict(gen_img)
74
+ # gen_img = Image.composite(gen_img, image.resize(gen_img.size), mask_image.resize(gen_img.size))
75
+ # out[i] = gen_img
76
 
77
  # return first generate PIL image
78
  json_imgs = {}
requirements.txt DELETED
@@ -1 +0,0 @@
1
- git+https://github.com/sberbank-ai/Real-ESRGAN.git