shahza1b commited on
Commit
5df167a
·
verified ·
1 Parent(s): 3e791eb

Update app.py

Browse files

Fix: Handle missing or empty 'layers' field in person_image

- Added validation to check if 'layers' exist and are non-empty before accessing the first layer in the submit function.
- If 'layers' are missing or empty, set the mask to None and auto-generate it using the AutoMasker.
- Ensured smooth fallback for cases where no manual mask is provided by automatically creating a mask.
- Enhanced stability by preventing 'IndexError' and ensuring graceful handling of input images without layers.

This commit improves error handling and ensures more robust functionality for generating virtual try-on results.

Files changed (1) hide show
  1. app.py +13 -13
app.py CHANGED
@@ -133,14 +133,19 @@ def submit_function(
133
  seed,
134
  show_type
135
  ):
136
- person_image, mask = person_image["background"], person_image["layers"][0]
137
- mask = Image.open(mask).convert("L")
138
- if len(np.unique(np.array(mask))) == 1:
139
- mask = None
 
 
 
 
 
 
140
  else:
141
- mask = np.array(mask)
142
- mask[mask > 0] = 255
143
- mask = Image.fromarray(mask)
144
 
145
  tmp_folder = args.output_dir
146
  date_str = datetime.now().strftime("%Y%m%d%H%M%S")
@@ -168,7 +173,6 @@ def submit_function(
168
  mask = mask_processor.blur(mask, blur_factor=9)
169
 
170
  # Inference
171
- # try:
172
  result_image = pipeline(
173
  image=person_image,
174
  condition_image=cloth_image,
@@ -177,10 +181,6 @@ def submit_function(
177
  guidance_scale=guidance_scale,
178
  generator=generator
179
  )[0]
180
- # except Exception as e:
181
- # raise gr.Error(
182
- # "An error occurred. Please try again later: {}".format(e)
183
- # )
184
 
185
  # Post-process
186
  masked_person = vis_mask(person_image, mask)
@@ -195,7 +195,7 @@ def submit_function(
195
  conditions = image_grid([person_image, cloth_image], 2, 1)
196
  else:
197
  condition_width = width // 3
198
- conditions = image_grid([person_image, masked_person , cloth_image], 3, 1)
199
  conditions = conditions.resize((condition_width, height), Image.NEAREST)
200
  new_result_image = Image.new("RGB", (width + condition_width + 5, height))
201
  new_result_image.paste(conditions, (0, 0))
 
133
  seed,
134
  show_type
135
  ):
136
+ # Check if layers exist and are not empty
137
+ if "layers" in person_image and person_image["layers"]:
138
+ person_image, mask = person_image["background"], person_image["layers"][0]
139
+ mask = Image.open(mask).convert("L")
140
+ if len(np.unique(np.array(mask))) == 1: # All mask values are the same (empty mask)
141
+ mask = None
142
+ else:
143
+ mask = np.array(mask)
144
+ mask[mask > 0] = 255 # Convert to binary mask (0 or 255)
145
+ mask = Image.fromarray(mask)
146
  else:
147
+ person_image = person_image["background"]
148
+ mask = None # No mask is provided, it will be auto-generated
 
149
 
150
  tmp_folder = args.output_dir
151
  date_str = datetime.now().strftime("%Y%m%d%H%M%S")
 
173
  mask = mask_processor.blur(mask, blur_factor=9)
174
 
175
  # Inference
 
176
  result_image = pipeline(
177
  image=person_image,
178
  condition_image=cloth_image,
 
181
  guidance_scale=guidance_scale,
182
  generator=generator
183
  )[0]
 
 
 
 
184
 
185
  # Post-process
186
  masked_person = vis_mask(person_image, mask)
 
195
  conditions = image_grid([person_image, cloth_image], 2, 1)
196
  else:
197
  condition_width = width // 3
198
+ conditions = image_grid([person_image, masked_person, cloth_image], 3, 1)
199
  conditions = conditions.resize((condition_width, height), Image.NEAREST)
200
  new_result_image = Image.new("RGB", (width + condition_width + 5, height))
201
  new_result_image.paste(conditions, (0, 0))