yiren98 commited on
Commit
98973cb
·
1 Parent(s): d1bd23e
Files changed (1) hide show
  1. gradio_app.py +6 -1
gradio_app.py CHANGED
@@ -2,7 +2,6 @@ import spaces
2
  import gradio as gr
3
  import torch
4
  import numpy as np
5
- import random
6
  from PIL import Image
7
  from accelerate import Accelerator
8
  import os
@@ -120,6 +119,12 @@ class ResizeWithPadding:
120
 
121
  width, height = img.size
122
 
 
 
 
 
 
 
123
  if width == height:
124
  img = img.resize((self.size, self.size), Image.LANCZOS)
125
  else:
 
2
  import gradio as gr
3
  import torch
4
  import numpy as np
 
5
  from PIL import Image
6
  from accelerate import Accelerator
7
  import os
 
119
 
120
  width, height = img.size
121
 
122
+ # Convert to RGB to remove transparency, fill with white background if necessary
123
+ if img.mode in ('RGBA', 'LA') or (img.mode == 'P' and 'transparency' in img.info):
124
+ background = Image.new("RGB", img.size, (fill, fill, fill))
125
+ background.paste(img, mask=img.split()[-1]) # Use alpha channel as mask
126
+ img = background
127
+
128
  if width == height:
129
  img = img.resize((self.size, self.size), Image.LANCZOS)
130
  else: