avans06 commited on
Commit
129b135
·
1 Parent(s): 645bf54

Fixed the handling behavior when face_model is set to None.

Browse files
Files changed (3) hide show
  1. README.md +1 -1
  2. app.py +57 -24
  3. requirements.txt +1 -1
README.md CHANGED
@@ -4,7 +4,7 @@ emoji: 📈
4
  colorFrom: blue
5
  colorTo: gray
6
  sdk: gradio
7
- sdk_version: 5.25.1
8
  app_file: app.py
9
  pinned: true
10
  license: apache-2.0
 
4
  colorFrom: blue
5
  colorTo: gray
6
  sdk: gradio
7
+ sdk_version: 5.35.0
8
  app_file: app.py
9
  pinned: true
10
  license: apache-2.0
app.py CHANGED
@@ -219,6 +219,20 @@ Phhofm: Restoration, 4x ESRGAN model for photography, trained using the Real-ESR
219
  Phhofm: Restoration, 4x ESRGAN model for photography, trained with realistic noise, lens blur, jpg and webp re-compression.
220
  ESRGAN version of 4xNomosWebPhoto_RealPLKSR, trained on the same dataset and in the same way."""],
221
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
222
  # DATNet
223
  "4xNomos8kDAT.pth" : ["https://github.com/Phhofm/models/releases/download/4xNomos8kDAT/4xNomos8kDAT.pth",
224
  "https://openmodeldb.info/models/4x-Nomos8kDAT",
@@ -431,7 +445,7 @@ example_list = ["images/a01.jpg", "images/a02.jpg", "images/a03.jpg", "images/a0
431
  def get_model_type(model_name):
432
  # Define model type mappings based on key parts of the model names
433
  model_type = "other"
434
- if any(value in model_name.lower() for value in ("4x-animesharp.pth", "sudo-realesrgan")):
435
  model_type = "ESRGAN"
436
  elif "srformer" in model_name.lower():
437
  model_type = "SRFormer"
@@ -726,7 +740,7 @@ class Upscale:
726
  return new_image
727
 
728
  def enhance(self, img, outscale=None):
729
- # img: numpy
730
  h_input, w_input = img.shape[0:2]
731
  pil_img = self.cv2pil(img)
732
  pil_img = self.__call__(pil_img)
@@ -791,17 +805,17 @@ class Upscale:
791
  progressRatio = 0.5 if upscale_model and face_restoration else 1
792
  print(f"progressRatio: {progressRatio}")
793
  current_progress = 0
794
- progress(0, desc="Initialize model start")
795
  if upscale_model:
796
  self.initBGUpscaleModel(upscale_model)
797
  current_progress += progressRatio/progressTotal;
798
- progress(current_progress, desc="Initialize BG upscale model finished")
799
  timer.checkpoint(f"Initialize BG upscale model")
800
 
801
  if face_restoration:
802
  self.initFaceEnhancerModel(face_restoration, face_detection)
803
  current_progress += progressRatio/progressTotal;
804
- progress(current_progress, desc="Initialize face enhancer model finished")
805
  timer.checkpoint(f"Initialize face enhancer model")
806
 
807
  timer.report()
@@ -825,30 +839,39 @@ class Upscale:
825
  # Dictionary to track counters for each filename
826
  name_counters = defaultdict(int)
827
  for gallery_idx, value in enumerate(gallery):
 
828
  try:
829
- img_path = str(value[0])
 
 
 
830
  img_name = os.path.basename(img_path)
831
  basename, extension = os.path.splitext(img_name)
832
- # Increment the counter for the current name
833
  name_counters[img_name] += 1
834
  if name_counters[img_name] > 1:
835
  basename = f"{basename}_{name_counters[img_name] - 1:02d}"
836
 
837
  img_cv2 = cv2.imdecode(np.fromfile(img_path, np.uint8), cv2.IMREAD_UNCHANGED) # numpy.ndarray
838
 
 
 
 
 
 
839
  img_mode = "RGBA" if len(img_cv2.shape) == 3 and img_cv2.shape[2] == 4 else None
840
  if len(img_cv2.shape) == 2: # for gray inputs
841
  img_cv2 = cv2.cvtColor(img_cv2, cv2.COLOR_GRAY2BGR)
842
- print(f"> image{gallery_idx:02d}, {img_cv2.shape}:")
843
 
844
  bg_upsample_img = None
845
- if self.realesrganer and hasattr(self.realesrganer, "enhance"):
846
  bg_upsample_img, _ = auto_split_upscale(img_cv2, self.realesrganer.enhance, self.scale) if is_auto_split_upscale else self.realesrganer.enhance(img_cv2, outscale=self.scale)
847
  current_progress += progressRatio/progressTotal;
848
- progress(current_progress, desc=f"image{gallery_idx:02d}, Background upscale Section")
849
- timer.checkpoint(f"image{gallery_idx:02d}, Background upscale Section")
850
 
851
- if self.face_enhancer:
852
  cropped_faces, restored_aligned, bg_upsample_img = self.face_enhancer.enhance(img_cv2, has_aligned=False, only_center_face=face_detection_only_center, paste_back=True, bg_upsample_img=bg_upsample_img, eye_dist_threshold=face_detection_threshold)
853
  # save faces
854
  if cropped_faces and restored_aligned:
@@ -871,11 +894,16 @@ class Upscale:
871
  files.append(save_restore_path)
872
  files.append(save_cmp_path)
873
  current_progress += progressRatio/progressTotal;
874
- progress(current_progress, desc=f"image{gallery_idx:02d}, Face enhancer Section")
875
- timer.checkpoint(f"image{gallery_idx:02d}, Face enhancer Section")
876
 
877
  restored_img = bg_upsample_img
878
- timer.report()
 
 
 
 
 
879
 
880
  if not extension:
881
  extension = ".png" if img_mode == "RGBA" else ".jpg" # RGBA images should be saved in png format
@@ -886,29 +914,34 @@ class Upscale:
886
  restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
887
  files.append(save_path)
888
  except RuntimeError as error:
 
 
 
 
889
  print(traceback.format_exc())
890
- print('Error', error)
 
891
 
892
- progress(1, desc=f"Execution completed")
893
- timer.report_all() # Print all recorded times
894
  # Close zip files
895
  zipf_cropf.close()
896
  zipf_restoref.close()
897
  zipf_cmp.close()
898
  zipf_restore.close()
899
  except Exception as error:
 
900
  print(traceback.format_exc())
901
- print("global exception: ", error)
902
  return None, None
903
  finally:
904
- if self.face_enhancer:
905
  self.face_enhancer._cleanup()
906
- else:
907
- # Free GPU memory and clean up resources
908
  torch.cuda.empty_cache()
909
- gc.collect()
910
 
911
- return files, [zip_cropf_path, zip_restoref_path, zip_cmp_path, zip_restore_path]
912
 
913
 
914
  def find_max_numbers(self, state_dict, findkeys):
 
219
  Phhofm: Restoration, 4x ESRGAN model for photography, trained with realistic noise, lens blur, jpg and webp re-compression.
220
  ESRGAN version of 4xNomosWebPhoto_RealPLKSR, trained on the same dataset and in the same way."""],
221
 
222
+
223
+ "4x_foolhardy_Remacri.pth": ["https://civitai.com/api/download/models/164821?type=Model&format=PickleTensor",
224
+ "https://openmodeldb.info/models/4x-Remacri",
225
+ """Original
226
+ FoolhardyVEVO: A creation of BSRGAN with more details and less smoothing, made by interpolating IRL models such as Siax,
227
+ Superscale, Superscale Artisoft, Pixel Perfect, etc. This was, things like skin and other details don't become mushy and blurry."""],
228
+
229
+ "4x_foolhardy_Remacri_ExtraSmoother.pth": ["https://civitai.com/api/download/models/164822?type=Model&format=PickleTensor",
230
+ "https://openmodeldb.info/models/4x-Remacri",
231
+ """ExtraSmoother
232
+ FoolhardyVEVO: A creation of BSRGAN with more details and less smoothing, made by interpolating IRL models such as Siax,
233
+ Superscale, Superscale Artisoft, Pixel Perfect, etc. This was, things like skin and other details don't become mushy and blurry."""],
234
+
235
+
236
  # DATNet
237
  "4xNomos8kDAT.pth" : ["https://github.com/Phhofm/models/releases/download/4xNomos8kDAT/4xNomos8kDAT.pth",
238
  "https://openmodeldb.info/models/4x-Nomos8kDAT",
 
445
  def get_model_type(model_name):
446
  # Define model type mappings based on key parts of the model names
447
  model_type = "other"
448
+ if any(value in model_name.lower() for value in ("4x-animesharp.pth", "sudo-realesrgan", "remacri")):
449
  model_type = "ESRGAN"
450
  elif "srformer" in model_name.lower():
451
  model_type = "SRFormer"
 
740
  return new_image
741
 
742
  def enhance(self, img, outscale=None):
743
+ # img: numpy array
744
  h_input, w_input = img.shape[0:2]
745
  pil_img = self.cv2pil(img)
746
  pil_img = self.__call__(pil_img)
 
805
  progressRatio = 0.5 if upscale_model and face_restoration else 1
806
  print(f"progressRatio: {progressRatio}")
807
  current_progress = 0
808
+ progress(0, desc="Initializing models...")
809
  if upscale_model:
810
  self.initBGUpscaleModel(upscale_model)
811
  current_progress += progressRatio/progressTotal;
812
+ progress(current_progress, desc="BG upscale model initialized.")
813
  timer.checkpoint(f"Initialize BG upscale model")
814
 
815
  if face_restoration:
816
  self.initFaceEnhancerModel(face_restoration, face_detection)
817
  current_progress += progressRatio/progressTotal;
818
+ progress(current_progress, desc="Face enhancer model initialized.")
819
  timer.checkpoint(f"Initialize face enhancer model")
820
 
821
  timer.report()
 
839
  # Dictionary to track counters for each filename
840
  name_counters = defaultdict(int)
841
  for gallery_idx, value in enumerate(gallery):
842
+ img_path = None
843
  try:
844
+ if value is None or not value:
845
+ print(f"Warning: Invalid gallery item at index {gallery_idx}. Skipping.")
846
+ continue
847
+ img_path = str(value[0]) # value is often a list/tuple like [filepath, caption]
848
  img_name = os.path.basename(img_path)
849
  basename, extension = os.path.splitext(img_name)
850
+ # Increment the counter for the current name if it appears multiple times
851
  name_counters[img_name] += 1
852
  if name_counters[img_name] > 1:
853
  basename = f"{basename}_{name_counters[img_name] - 1:02d}"
854
 
855
  img_cv2 = cv2.imdecode(np.fromfile(img_path, np.uint8), cv2.IMREAD_UNCHANGED) # numpy.ndarray
856
 
857
+ if img_cv2 is None:
858
+ print(f"Warning: Could not read or decode image '{img_path}'. Skipping this image.")
859
+ # Skip this iteration and process the next image
860
+ continue
861
+
862
  img_mode = "RGBA" if len(img_cv2.shape) == 3 and img_cv2.shape[2] == 4 else None
863
  if len(img_cv2.shape) == 2: # for gray inputs
864
  img_cv2 = cv2.cvtColor(img_cv2, cv2.COLOR_GRAY2BGR)
865
+ print(f"> Processing image {gallery_idx:02d}, Shape: {img_cv2.shape}")
866
 
867
  bg_upsample_img = None
868
+ if upscale_model and self.realesrganer and hasattr(self.realesrganer, "enhance"):
869
  bg_upsample_img, _ = auto_split_upscale(img_cv2, self.realesrganer.enhance, self.scale) if is_auto_split_upscale else self.realesrganer.enhance(img_cv2, outscale=self.scale)
870
  current_progress += progressRatio/progressTotal;
871
+ progress(current_progress, desc=f"Image {gallery_idx:02d}: Background upscaling...")
872
+ timer.checkpoint(f"Image {gallery_idx:02d}: Background upscale section")
873
 
874
+ if face_restoration and self.face_enhancer:
875
  cropped_faces, restored_aligned, bg_upsample_img = self.face_enhancer.enhance(img_cv2, has_aligned=False, only_center_face=face_detection_only_center, paste_back=True, bg_upsample_img=bg_upsample_img, eye_dist_threshold=face_detection_threshold)
876
  # save faces
877
  if cropped_faces and restored_aligned:
 
894
  files.append(save_restore_path)
895
  files.append(save_cmp_path)
896
  current_progress += progressRatio/progressTotal;
897
+ progress(current_progress, desc=f"Image {gallery_idx:02d}: Face enhancement...")
898
+ timer.checkpoint(f"Image {gallery_idx:02d}: Face enhancer section")
899
 
900
  restored_img = bg_upsample_img
901
+ timer.report() # Report time for this image
902
+
903
+ # Handle cases where image processing might still result in None
904
+ if restored_img is None:
905
+ print(f"Warning: Processing resulted in no image for '{img_path}'. Skipping output.")
906
+ continue
907
 
908
  if not extension:
909
  extension = ".png" if img_mode == "RGBA" else ".jpg" # RGBA images should be saved in png format
 
914
  restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
915
  files.append(save_path)
916
  except RuntimeError as error:
917
+ print(f"Runtime Error while processing image {gallery_idx} ({img_path or 'unknown path'}): {error}")
918
+ print(traceback.format_exc())
919
+ except Exception as general_error:
920
+ print(f"An unexpected error occurred while processing image {gallery_idx} ({img_path or 'unknown path'}): {general_error}")
921
  print(traceback.format_exc())
922
+ # Can still choose to continue to process the next image
923
+ continue
924
 
925
+ progress(1, desc=f"Processing complete.")
926
+ timer.report_all() # Print all recorded times for the whole batch
927
  # Close zip files
928
  zipf_cropf.close()
929
  zipf_restoref.close()
930
  zipf_cmp.close()
931
  zipf_restore.close()
932
  except Exception as error:
933
+ print(f"Global exception occurred: {error}")
934
  print(traceback.format_exc())
 
935
  return None, None
936
  finally:
937
+ if hasattr(self, 'face_enhancer') and self.face_enhancer:
938
  self.face_enhancer._cleanup()
939
+ # Free GPU memory and clean up resources
940
+ if torch.cuda.is_available():
941
  torch.cuda.empty_cache()
942
+ gc.collect()
943
 
944
+ return files, [zip_cropf_path, zip_restoref_path, zip_cmp_path, zip_restore_path] if face_restoration else [zip_restore_path]
945
 
946
 
947
  def find_max_numbers(self, state_dict, findkeys):
requirements.txt CHANGED
@@ -1,6 +1,6 @@
1
  --extra-index-url https://download.pytorch.org/whl/cu124
2
 
3
- gradio==5.25.1
4
 
5
  basicsr @ git+https://github.com/avan06/BasicSR
6
  facexlib @ git+https://github.com/avan06/facexlib
 
1
  --extra-index-url https://download.pytorch.org/whl/cu124
2
 
3
+ gradio==5.35.0
4
 
5
  basicsr @ git+https://github.com/avan06/BasicSR
6
  facexlib @ git+https://github.com/avan06/facexlib