ohayonguy commited on
Commit
2cd7eda
1 Parent(s): 1fef40b

trying to fix interface

Browse files
Files changed (1) hide show
  1. app.py +7 -13
app.py CHANGED
@@ -52,12 +52,10 @@ def generate_reconstructions(pmrf_model, x, y, non_noisy_z0, num_flow_steps, dev
52
  dt = (1.0 / num_flow_steps) * (1.0 - pmrf_model.hparams.eps)
53
  x_t_next = source_dist_samples.clone()
54
  t_one = torch.ones(x.shape[0], device=device)
55
- pbar = tqdm(range(num_flow_steps))
56
- for i in pbar:
57
  num_t = (i / num_flow_steps) * (1.0 - pmrf_model.hparams.eps) + pmrf_model.hparams.eps
58
  v_t_next = pmrf_model(x_t=x_t_next, t=t_one * num_t, y=y).to(x_t_next.dtype)
59
  x_t_next = x_t_next.clone() + v_t_next * dt
60
- pbar.set_description(f'Flow step {i}')
61
 
62
  return x_t_next.clip(0, 1).to(torch.float32)
63
 
@@ -76,8 +74,9 @@ def enhance_face(img, face_helper, has_aligned, num_flow_steps, only_center_face
76
  # TODO: even with eye_dist_threshold, it will still introduce wrong detections and restorations.
77
  # align and warp each face
78
  face_helper.align_warp_face()
 
79
  # face restoration
80
- for i, cropped_face in enumerate(face_helper.cropped_faces):
81
  # prepare data
82
  h, w = cropped_face.shape[0], cropped_face.shape[1]
83
  cropped_face = cv2.resize(cropped_face, (512, 512), interpolation=cv2.INTER_LINEAR)
@@ -114,8 +113,7 @@ def enhance_face(img, face_helper, has_aligned, num_flow_steps, only_center_face
114
  def inference(seed, randomize_seed, img, aligned, scale, num_flow_steps,
115
  progress=gr.Progress(track_tqdm=True)):
116
  if img is None:
117
- gr.Info("Please upload an image before submitting")
118
- return [None, None, None]
119
  if randomize_seed:
120
  seed = random.randint(0, MAX_SEED)
121
  torch.manual_seed(seed)
@@ -127,8 +125,7 @@ def inference(seed, randomize_seed, img, aligned, scale, num_flow_steps,
127
 
128
  h, w = img.shape[0:2]
129
  if h > 4500 or w > 4500:
130
- print('Image size too large.')
131
- return None, None
132
 
133
  face_helper = FaceRestoreHelper(
134
  scale,
@@ -172,11 +169,9 @@ You may use this demo to enhance the quality of any image which contains faces.
172
 
173
  Please refer to our project's page for more details: https://pmrf-ml.github.io/.
174
 
175
- ---
176
-
177
  *Notes* :
178
 
179
- 1. Our model is designed to restore aligned face images, but here we incorporate mechanisms that allow restoring the quality of any image that contains any number of faces. Thus, the resulting quality of such general images is not guaranteed.
180
  2. Images that are too large won't work due to memory constraints.
181
 
182
  ---
@@ -201,8 +196,7 @@ If you find our work useful, please help to ⭐ our <a href='https://github.com/
201
 
202
  📋 **License**
203
 
204
- This project is released under the <a rel="license" href="https://github.com/ohayonguy/PMRF/blob/master/LICENSE">MIT license</a>.
205
- Redistribution and use for non-commercial purposes should follow this license.
206
 
207
  📧 **Contact**
208
 
 
52
  dt = (1.0 / num_flow_steps) * (1.0 - pmrf_model.hparams.eps)
53
  x_t_next = source_dist_samples.clone()
54
  t_one = torch.ones(x.shape[0], device=device)
55
+ for i in tqdm(range(num_flow_steps)):
 
56
  num_t = (i / num_flow_steps) * (1.0 - pmrf_model.hparams.eps) + pmrf_model.hparams.eps
57
  v_t_next = pmrf_model(x_t=x_t_next, t=t_one * num_t, y=y).to(x_t_next.dtype)
58
  x_t_next = x_t_next.clone() + v_t_next * dt
 
59
 
60
  return x_t_next.clip(0, 1).to(torch.float32)
61
 
 
74
  # TODO: even with eye_dist_threshold, it will still introduce wrong detections and restorations.
75
  # align and warp each face
76
  face_helper.align_warp_face()
77
+ gr.Info(f"Identified {len(face_helper.cropped_faces)} faces in the image. The algorithm will enhance the quality of each face.")
78
  # face restoration
79
+ for i, cropped_face in tqdm(enumerate(face_helper.cropped_faces)):
80
  # prepare data
81
  h, w = cropped_face.shape[0], cropped_face.shape[1]
82
  cropped_face = cv2.resize(cropped_face, (512, 512), interpolation=cv2.INTER_LINEAR)
 
113
  def inference(seed, randomize_seed, img, aligned, scale, num_flow_steps,
114
  progress=gr.Progress(track_tqdm=True)):
115
  if img is None:
116
+ raise gr.Error("Please upload an image before submitting.")
 
117
  if randomize_seed:
118
  seed = random.randint(0, MAX_SEED)
119
  torch.manual_seed(seed)
 
125
 
126
  h, w = img.shape[0:2]
127
  if h > 4500 or w > 4500:
128
+ raise gr.Error('Image size too large.')
 
129
 
130
  face_helper = FaceRestoreHelper(
131
  scale,
 
169
 
170
  Please refer to our project's page for more details: https://pmrf-ml.github.io/.
171
 
 
 
172
  *Notes* :
173
 
174
+ 1. Our model is designed to restore aligned face images, where there is *only one* face in the image, and the face is centered. Here, however, we incorporate mechanisms that allow restoring the quality of *any* image that contains *any* number of faces. Thus, the resulting quality of such general images is not guaranteed.
175
  2. Images that are too large won't work due to memory constraints.
176
 
177
  ---
 
196
 
197
  📋 **License**
198
 
199
+ This project is released under the <a rel="license" href="https://github.com/ohayonguy/PMRF/blob/master/LICENSE">MIT license</a>.
 
200
 
201
  📧 **Contact**
202