Update app.py

app.py

@@ -29,14 +29,14 @@ SAFE_IMG_SIZE = 512
 SOURCE_AGE = 20
 TARGET_AGE = 80
 WRINKLE_STRENGTH = 0.42
-CONTRAST_BOOST = 1.10
-SHARPNESS_BOOST = 1.20
-ALPHA_HAIR = 0.95
+CONTRAST_BOOST = 1.15          # thoda badhaya
+SHARPNESS_BOOST = 1.35         # thoda badhaya
+ALPHA_HAIR = 0.92
 BLUR_RADIUS = 7
 EDGE_SMOOTHING = True
 USE_GFPGAN = True
 GFPGAN_UPSCALE = 1
-GFPGAN_WEIGHT = 0.5
+GFPGAN_WEIGHT = 0.75           # ← Quality ke liye badhaya (0.5 se 0.75)
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"🚀 Device: {DEVICE}")
@@ -88,11 +88,10 @@ def load_gfpgan():
             channel_multiplier=2,
             bg_upsampler=None,
             device=DEVICE,
-            half=False          # Safe setting for HF Space
+            half=False
         )
         print("✅ GFPGAN loaded successfully with FP32!")
         return gfpgan_restorer
-
     except Exception as e:
         print(f"❌ GFPGAN load failed: {e}")
         return None
@@ -104,6 +103,7 @@ def load_aging_model():
         return age_model
 
     print("Loading UNet aging model...")
+
     class DownLayer(nn.Module):
         def __init__(self, in_ch, out_ch):
             super().__init__()
@@ -239,6 +239,7 @@ def get_lips_mask(pil_image: Image.Image) -> np.ndarray:
                     return lips_mask
     return np.zeros((h, w), dtype=np.float32)
 
+
 def exclude_lips_from_mask(beard_mask: np.ndarray, pil_image: Image.Image) -> np.ndarray:
     if np.sum(beard_mask) == 0:
         return beard_mask
@@ -250,6 +251,7 @@ def exclude_lips_from_mask(beard_mask: np.ndarray, pil_image: Image.Image) -> np
     beard_mask = cv2.GaussianBlur(beard_mask, (5, 5), 1)
     return beard_mask
 
+
 def get_beard_mask(pil_image: Image.Image) -> np.ndarray:
     temp_path = "temp_input.jpg"
     try:
@@ -281,6 +283,7 @@ def get_beard_mask(pil_image: Image.Image) -> np.ndarray:
         if os.path.exists(temp_path):
             os.remove(temp_path)
 
+
 def clean_mask(mask, min_area=150):
     mask = mask.astype(np.uint8)
     labeled, num = label(mask)
@@ -290,6 +293,7 @@ def clean_mask(mask, min_area=150):
             new_mask[labeled == i] = 1
     return new_mask
 
+
 def get_hair_mask_segformer(pil_image: Image.Image) -> np.ndarray:
     processor, parser = load_face_parser()
     inputs = processor(images=pil_image, return_tensors="pt").to(DEVICE)
@@ -313,6 +317,7 @@ def get_hair_mask_segformer(pil_image: Image.Image) -> np.ndarray:
     hair_mask = np.clip(hair_mask, 0, 1)
     return hair_mask
 
+
 def apply_hair_and_beard_color(image: Image.Image, hair_mask: np.ndarray, beard_mask: np.ndarray):
     combined_mask = np.maximum(hair_mask, beard_mask)
     if np.sum(combined_mask) == 0:
@@ -332,24 +337,27 @@ def apply_hair_and_beard_color(image: Image.Image, hair_mask: np.ndarray, beard_
     result = (1 - alpha * combined_mask[..., np.newaxis]) * img_np + (alpha * combined_mask[..., np.newaxis]) * white_layer
     result = np.clip(result, 0, 255).astype(np.uint8)
     result_pil = Image.fromarray(result)
-    result_pil = result_pil.filter(ImageFilter.UnsharpMask(1.2, 140, 2))
+    result_pil = result_pil.filter(ImageFilter.UnsharpMask(1.3, 150, 2))
     return result_pil
 
+
 def post_correct_aged(original: Image.Image, aged: Image.Image) -> Image.Image:
     orig_np = np.array(original)
     aged_np = np.array(aged)
     matched = match_histograms(aged_np, orig_np, channel_axis=-1)
     matched_img = Image.fromarray(np.clip(matched, 0, 255).astype(np.uint8))
-    matched_img = ImageEnhance.Brightness(matched_img).enhance(1.10)
-    matched_img = ImageEnhance.Contrast(matched_img).enhance(1.06)
+    matched_img = ImageEnhance.Brightness(matched_img).enhance(1.12)
+    matched_img = ImageEnhance.Contrast(matched_img).enhance(1.18)
     return matched_img
 
+
 def enhance_texture(img: Image.Image) -> Image.Image:
-    img = img.filter(ImageFilter.UnsharpMask(2, 160, 3))
+    img = img.filter(ImageFilter.UnsharpMask(radius=2.5, percent=180, threshold=3))  # Stronger sharpening
     img = ImageEnhance.Contrast(img).enhance(CONTRAST_BOOST)
     img = ImageEnhance.Sharpness(img).enhance(SHARPNESS_BOOST)
     return img
 
+
 # ================== MAIN PROCESSING FUNCTION ==================
 def process_face_aging(input_image: Image.Image) -> Image.Image:
     if input_image is None:
@@ -374,8 +382,9 @@ def process_face_aging(input_image: Image.Image) -> Image.Image:
             blended = (1 - alpha) * rgb_tensor.unsqueeze(0) + alpha * raw_output
             blended = blended.clamp(0, 1).squeeze(0)
             final_aged = TF.to_pil_image(blended).resize((ow, oh), Image.LANCZOS)
-            final_aged = enhance_texture(final_aged)
-            final_aged = post_correct_aged(orig, final_aged)
+
+        final_aged = enhance_texture(final_aged)
+        final_aged = post_correct_aged(orig, final_aged)
 
         print(" Generating hair mask...")
         hair_mask = get_hair_mask_segformer(final_aged)
@@ -394,8 +403,11 @@ def process_face_aging(input_image: Image.Image) -> Image.Image:
                     img_cv = cv2.cvtColor(np.array(final_img), cv2.COLOR_RGB2BGR)
                     with torch.cuda.amp.autocast(enabled=(DEVICE.type == "cuda")):
                         _, _, restored_cv = gfpgan.enhance(
-                            img_cv, has_aligned=False, only_center_face=False, 
-                            paste_back=True, weight=GFPGAN_WEIGHT
+                            img_cv, 
+                            has_aligned=False, 
+                            only_center_face=False, 
+                            paste_back=True, 
+                            weight=GFPGAN_WEIGHT
                         )
                     final_img = Image.fromarray(cv2.cvtColor(restored_cv, cv2.COLOR_BGR2RGB))
                 except Exception as e:
@@ -412,6 +424,7 @@ def process_face_aging(input_image: Image.Image) -> Image.Image:
         traceback.print_exc()
         raise
 
+
 # ================== FASTAPI SETUP ==================
 app = FastAPI(title="Face Aging + White Hair & Beard API")
 
@@ -448,6 +461,7 @@ async def age_face(file: UploadFile = File(...)):
         if DEVICE.type == "cuda":
             torch.cuda.empty_cache()
 
+
 # For local testing
 if __name__ == "__main__":
     import uvicorn