Update app.py

app.py

@@ -9,7 +9,7 @@ import os
 import cv2
 import torch
 import torch.nn.functional as F
-import uuid, threading, time
+import uuid, threading, time, glob
 import gradio as gr
 
 from torchvision.transforms.functional import normalize
@@ -67,13 +67,16 @@ def imread(img_path):
     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     return img
 
-def delayed_remove(path, delay=90):
+def delayed_remove(path, delay=60):
     time.sleep(delay)
     try:
         if os.path.exists(path):
             os.remove(path)
+            print(f"[CLEANUP] removed: {path}")
+        else:
+            print(f"[CLEANUP] already gone: {path}")
     except Exception as e:
-        print(f"[WARN] Failed to remove {path}: {e}")
+        print(f"[CLEANUP] failed: {path} | {e}")
 
 # set enhancer with RealESRGAN
 def set_realesrgan():
@@ -176,6 +179,13 @@ def inference(image, face_align, background_enhance, face_upsample, upscale, cod
             # align and warp each face
             face_helper.align_warp_face()
 
+        if min(img.shape[:2]) > 1000 and num_det_faces > 15:
+            raise gr.Error(
+                "Too many faces detected (>15) in a high-resolution image. "
+                "To keep the demo responsive, this case is skipped. "
+                "For such inputs, please deploy the demo locally without this limit."
+            )
+            
         # face restoration for each cropped face
         for idx, cropped_face in enumerate(face_helper.cropped_faces):
             # prepare data
@@ -232,13 +242,14 @@ def inference(image, face_align, background_enhance, face_upsample, upscale, cod
         #save restored img
         save_path = f"output/{uuid.uuid4().hex}.png"
         imwrite(restored_img, save_path)
+        print(f"[SAVE] {save_path} | OUTPUT Number: {len(glob.glob("output/*.png"))}")
         
         threading.Thread(
             target=delayed_remove,
-            args=(save_path,10),
+            args=(save_path,30),
             daemon=True
         ).start()
-
+    
         return save_path
         
     except Exception as error:
@@ -314,7 +325,8 @@ demo = gr.Interface(
         gr.Number(value=2, label="Rescaling_Factor (up to 4)"),
         gr.Slider(0, 1, value=0.5, step=0.01, label='Codeformer_Fidelity (0 for better quality, 1 for better identity)')
     ], [
-        gr.Image(type="filepath", label="Output")
+        gr.Image(type="filepath", label="Output"),
+        gr.Markdown("Please download the output within 30 seconds.")
     ],
     title=title,
     description=description,