Kkordik commited on
Commit
dc12a0a
1 Parent(s): a998593

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +18 -49
app.py CHANGED
@@ -12,34 +12,27 @@ os.system("pip install gradio==4.29.0")
12
 
13
  login(token=os.getenv("HF_TOKEN"))
14
 
15
- # vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
16
- # pipeline = StableDiffusion3InpaintPipeline(vae=vae, torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("cuda")
17
-
18
  pipeline = StableDiffusion3InpaintPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
19
 
20
  def get_select_index(evt: gr.SelectData):
21
  return evt.index
22
-
23
  @spaces.GPU()
24
  def squarify_image(img):
25
- if(img.height > img.width): bg_size = img.height
26
- else: bg_size = img.width
27
- bg = Image.new(mode="RGB", size=(bg_size,bg_size), color="white")
28
- bg.paste(img, ( int((bg.width - bg.width)/2), 0) )
29
-
 
30
  return bg
31
 
32
  @spaces.GPU()
33
  def divisible_by_8(image):
34
  width, height = image.size
35
-
36
- # Calculate the new width and height that are divisible by 8
37
  new_width = (width // 8) * 8
38
  new_height = (height // 8) * 8
39
-
40
- # Resize the image
41
  resized_image = image.resize((new_width, new_height))
42
-
43
  return resized_image
44
 
45
  @spaces.GPU()
@@ -55,53 +48,29 @@ def clear_all():
55
  def generate(image_editor, prompt, neg_prompt, versions):
56
  start = time.time()
57
  image = image_editor['background'].convert('RGB')
58
-
59
- # Resize image
60
  image.thumbnail((1024, 1024))
61
  image = divisible_by_8(image)
62
  original_image_size = image.size
63
-
64
- # Mask layer
65
  layer = image_editor["layers"][0].resize(image.size)
66
-
67
- # Make image a square
68
  image = squarify_image(image)
69
-
70
- # Make sure mask is white with a black background
71
- mask = Image.new("RGBA", image.size, "WHITE")
72
  mask.paste(layer, (0, 0), layer)
73
  mask = ImageOps.invert(mask.convert('L'))
74
-
75
- # Inpaint
76
  pipeline.to("cuda")
77
- final_image = pipeline(prompt=prompt,
78
- image=image,
79
- mask_image=mask).images[0]
80
-
81
-
82
- # Make sure the longest side of image is 1024
83
- if (original_image_size[0] > original_image_size[1]):
84
- original_image_size = ( original_image_size[0] * (1024/original_image_size[0]) , original_image_size[1] * (1024/original_image_size[0]))
85
  else:
86
  original_image_size = (original_image_size[0] * (1024/original_image_size[1]), original_image_size[1] * (1024/original_image_size[1]))
87
-
88
-
89
- # Crop image to original aspect ratio
90
  final_image = final_image.crop((0, 0, original_image_size[0], original_image_size[1]))
91
-
92
- # gradio.ImageEditor requires a diction
93
  final_dict = {'background': final_image, 'layers': None, 'composite': final_image}
94
-
95
- # Add generated image to version gallery
96
- if(versions==None):
97
- final_gallery = [image_editor['background'] ,final_image]
98
- else:
99
  final_gallery = versions
100
  final_gallery.append(final_image)
101
-
102
  end = time.time()
103
  print('time:', end - start)
104
-
105
  return final_dict, gr.Gallery(value=final_gallery, visible=True), gr.update(visible=True), gr.update(visible=True)
106
 
107
  with gr.Blocks() as demo:
@@ -109,7 +78,7 @@ with gr.Blocks() as demo:
109
  # Inpainting SDXL Sketch Pad
110
  by [Tony Assi](https://www.tonyassi.com/)
111
 
112
- Please ❤️ this Space. I build custom AI apps for companies. <a href="mailto: tony.assi.media@gmail.com">Email me</a> for business inquiries.
113
  """)
114
 
115
  with gr.Row():
@@ -126,7 +95,7 @@ with gr.Blocks() as demo:
126
  selected = gr.Number(show_label=False, visible=False)
127
 
128
  gr.Examples(
129
- [[{'background':'./tony.jpg', 'layers':['./tony-mask.jpg'], 'composite':'./tony.jpg'}, 'black and white tuxedo, bowtie', 'ugly', None]],
130
  [sketch_pad, prompt, neg_prompt, version_gallery],
131
  [sketch_pad, version_gallery, restore_button, clear_button],
132
  generate,
@@ -134,8 +103,8 @@ with gr.Blocks() as demo:
134
  )
135
 
136
  version_gallery.select(get_select_index, None, selected)
137
- generate_button.click(fn=generate, inputs=[sketch_pad,prompt, neg_prompt, version_gallery], outputs=[sketch_pad, version_gallery, restore_button, clear_button])
138
  restore_button.click(fn=restore_version, inputs=[selected, version_gallery], outputs=sketch_pad)
139
  clear_button.click(clear_all, inputs=None, outputs=[sketch_pad, prompt, version_gallery, restore_button, clear_button])
140
 
141
- demo.launch()
 
12
 
13
  login(token=os.getenv("HF_TOKEN"))
14
 
 
 
 
15
  pipeline = StableDiffusion3InpaintPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
16
 
17
  def get_select_index(evt: gr.SelectData):
18
  return evt.index
19
+
20
  @spaces.GPU()
21
  def squarify_image(img):
22
+ if img.height > img.width:
23
+ bg_size = img.height
24
+ else:
25
+ bg_size = img.width
26
+ bg = Image.new(mode="RGB", size=(bg_size, bg_size), color="white")
27
+ bg.paste(img, (int((bg.width - bg.width) / 2), 0))
28
  return bg
29
 
30
  @spaces.GPU()
31
  def divisible_by_8(image):
32
  width, height = image.size
 
 
33
  new_width = (width // 8) * 8
34
  new_height = (height // 8) * 8
 
 
35
  resized_image = image.resize((new_width, new_height))
 
36
  return resized_image
37
 
38
  @spaces.GPU()
 
48
  def generate(image_editor, prompt, neg_prompt, versions):
49
  start = time.time()
50
  image = image_editor['background'].convert('RGB')
 
 
51
  image.thumbnail((1024, 1024))
52
  image = divisible_by_8(image)
53
  original_image_size = image.size
 
 
54
  layer = image_editor["layers"][0].resize(image.size)
 
 
55
  image = squarify_image(image)
56
+ mask = Image.new("RGBA", image.size, "WHITE")
 
 
57
  mask.paste(layer, (0, 0), layer)
58
  mask = ImageOps.invert(mask.convert('L'))
 
 
59
  pipeline.to("cuda")
60
+ final_image = pipeline(prompt=prompt, image=image, mask_image=mask).images[0]
61
+ if original_image_size[0] > original_image_size[1]:
62
+ original_image_size = (original_image_size[0] * (1024/original_image_size[0]), original_image_size[1] * (1024/original_image_size[0]))
 
 
 
 
 
63
  else:
64
  original_image_size = (original_image_size[0] * (1024/original_image_size[1]), original_image_size[1] * (1024/original_image_size[1]))
 
 
 
65
  final_image = final_image.crop((0, 0, original_image_size[0], original_image_size[1]))
 
 
66
  final_dict = {'background': final_image, 'layers': None, 'composite': final_image}
67
+ if versions is None:
68
+ final_gallery = [image_editor['background'], final_image]
69
+ else:
 
 
70
  final_gallery = versions
71
  final_gallery.append(final_image)
 
72
  end = time.time()
73
  print('time:', end - start)
 
74
  return final_dict, gr.Gallery(value=final_gallery, visible=True), gr.update(visible=True), gr.update(visible=True)
75
 
76
  with gr.Blocks() as demo:
 
78
  # Inpainting SDXL Sketch Pad
79
  by [Tony Assi](https://www.tonyassi.com/)
80
 
81
+ Please ❤️ this Space. I build custom AI apps for companies. <a href="mailto:tony.assi.media@gmail.com">Email me</a> for business inquiries.
82
  """)
83
 
84
  with gr.Row():
 
95
  selected = gr.Number(show_label=False, visible=False)
96
 
97
  gr.Examples(
98
+ [[{'background':'./tony.jpg', 'layers':['./tony-mask.jpg'], 'composite':'./tony.jpg'}, 'black and white tuxedo, bowtie', 'ugly', None]],
99
  [sketch_pad, prompt, neg_prompt, version_gallery],
100
  [sketch_pad, version_gallery, restore_button, clear_button],
101
  generate,
 
103
  )
104
 
105
  version_gallery.select(get_select_index, None, selected)
106
+ generate_button.click(fn=generate, inputs=[sketch_pad, prompt, neg_prompt, version_gallery], outputs=[sketch_pad, version_gallery, restore_button, clear_button])
107
  restore_button.click(fn=restore_version, inputs=[selected, version_gallery], outputs=sketch_pad)
108
  clear_button.click(clear_all, inputs=None, outputs=[sketch_pad, prompt, version_gallery, restore_button, clear_button])
109
 
110
+ demo.launch()