tonyassi commited on
Commit
4ed7a63
1 Parent(s): 3bec293

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -4
app.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import gradio as gr
2
  from diffusers import AutoPipelineForInpainting, AutoencoderKL
3
  import torch
@@ -6,9 +7,11 @@ from PIL import Image, ImageOps
6
  vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
7
  pipeline = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", vae=vae, torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("cuda")
8
 
 
9
  def get_select_index(evt: gr.SelectData):
10
  return evt.index
11
-
 
12
  def squarify_image(img):
13
  if(img.height > img.width): bg_size = img.height
14
  else: bg_size = img.width
@@ -16,7 +19,8 @@ def squarify_image(img):
16
  bg.paste(img, ( int((bg.width - bg.width)/2), 0) )
17
 
18
  return bg
19
-
 
20
  def divisible_by_8(image):
21
  width, height = image.size
22
 
@@ -29,12 +33,12 @@ def divisible_by_8(image):
29
 
30
  return resized_image
31
 
32
-
33
  def restore_version(index, versions):
34
  final_dict = {'background': versions[index][0], 'layers': None, 'composite': versions[index][0]}
35
  return final_dict
36
 
37
-
38
  def generate(image_editor, prompt, neg_prompt, versions):
39
  image = image_editor['background'].convert('RGB')
40
 
@@ -55,6 +59,7 @@ def generate(image_editor, prompt, neg_prompt, versions):
55
  mask = ImageOps.invert(mask.convert('L'))
56
 
57
  # Inpaint
 
58
  final_image = pipeline(prompt=prompt,
59
  image=image,
60
  mask_image=mask).images[0]
 
1
+ import spaces
2
  import gradio as gr
3
  from diffusers import AutoPipelineForInpainting, AutoencoderKL
4
  import torch
 
7
  vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
8
  pipeline = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", vae=vae, torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("cuda")
9
 
10
+ @spaces.GPU()
11
  def get_select_index(evt: gr.SelectData):
12
  return evt.index
13
+
14
+ @spaces.GPU()
15
  def squarify_image(img):
16
  if(img.height > img.width): bg_size = img.height
17
  else: bg_size = img.width
 
19
  bg.paste(img, ( int((bg.width - bg.width)/2), 0) )
20
 
21
  return bg
22
+
23
+ @spaces.GPU()
24
  def divisible_by_8(image):
25
  width, height = image.size
26
 
 
33
 
34
  return resized_image
35
 
36
+ @spaces.GPU()
37
  def restore_version(index, versions):
38
  final_dict = {'background': versions[index][0], 'layers': None, 'composite': versions[index][0]}
39
  return final_dict
40
 
41
+ @spaces.GPU()
42
  def generate(image_editor, prompt, neg_prompt, versions):
43
  image = image_editor['background'].convert('RGB')
44
 
 
59
  mask = ImageOps.invert(mask.convert('L'))
60
 
61
  # Inpaint
62
+ pipeline.to("cuda")
63
  final_image = pipeline(prompt=prompt,
64
  image=image,
65
  mask_image=mask).images[0]