Update handler.py

handler.py

@@ -1,31 +1,36 @@
 import os
-from diffusers import DiffusionPipeline
+import torch
+from gfpgan import GFPGANer
 from PIL import Image
+import cv2
 
 class EndpointHandler:
-    def __init__(self, model_dir="stabilityai/stable-diffusion-x4-upscaler"):
-        # Load the diffusion pipeline
-        self.pipeline = DiffusionPipeline.from_pretrained(model_dir)
-
+    def __init__(self, model_path='GFPGANv1.4.pth'):
+        # Load the GFPGAN model
+        self.model_path = model_path
+        self.bg_upsampler = None  # You can set this to RealESRGANer if needed
+        self.face_enhancer = GFPGANer(
+            model_path=self.model_path, upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=self.bg_upsampler)
+        
         # Ensure the output directory exists
         os.makedirs('output', exist_ok=True)
 
-    def upscale_image(self, image_path):
+    def enhance_image(self, image_path):
         # Load the image
-        image = Image.open(image_path)
+        img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
         
-        # Upscale the image
-        upscaled_image = self.pipeline(image)
+        # Perform face enhancement
+        _, _, output = self.face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
         
-        # Save the upscaled image
-        save_path = "output/upscaled_image.png"
-        upscaled_image.save(save_path)
-        
-        return upscaled_image, save_path
+        # Save the enhanced image
+        save_path = "output/enhanced_image.png"
+        cv2.imwrite(save_path, output)
+
+        return output, save_path
 
 # Example usage for testing
 if __name__ == "__main__":
     handler = EndpointHandler()
     test_image_path = 'path_to_test_image.jpg'  # Replace with your test image path
-    upscaled_image, save_path = handler.upscale_image(test_image_path)
-    print(f"Upscaled image saved at {save_path}")
+    enhanced_image, save_path = handler.enhance_image(test_image_path)
+    print(f"Enhanced image saved at {save_path}")