Fixed the handling behavior when face_model is set to None.
Browse files- README.md +1 -1
- app.py +57 -24
- 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.
|
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="
|
795 |
if upscale_model:
|
796 |
self.initBGUpscaleModel(upscale_model)
|
797 |
current_progress += progressRatio/progressTotal;
|
798 |
-
progress(current_progress, desc="
|
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="
|
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 |
-
|
|
|
|
|
|
|
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"
|
849 |
-
timer.checkpoint(f"
|
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"
|
875 |
-
timer.checkpoint(f"
|
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 |
-
|
|
|
891 |
|
892 |
-
progress(1, desc=f"
|
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 |
-
|
907 |
-
|
908 |
torch.cuda.empty_cache()
|
909 |
-
|
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.
|
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
|