Update inference_realesrgan.py

inference_realesrgan.py

@@ -4,8 +4,9 @@ import torch
 import cv2
 import numpy as np
 from PIL import Image
+from realesrgan import RealESRGANer
 
-# 🚀 Load Real-ESRGAN model manually from weight file
+# Load Real-ESRGAN model manually from weight file
 MODEL_PATH = "weight/RealESRGAN_x4plus.pth"
 
 # Ensure the model exists
@@ -13,21 +14,13 @@ if not os.path.exists(MODEL_PATH):
     print(f"[ERROR] Model file {MODEL_PATH} not found. Please download it using download_model.py.")
     exit(1)
 
-# 🚀 Fix: Manually load the ESRGAN model without requiring 'realesrgan' package
-from torch.hub import load_state_dict_from_url
-model_url = "https://huggingface.co/va1/CViT-Deepfake-Model/resolve/main/RealESRGAN_x4plus.pth"
-
-# Load weights if missing
-if not os.path.exists(MODEL_PATH):
-    torch.hub.download_url_to_file(model_url, MODEL_PATH)
-
-# Initialize ESRGAN model
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# 🚀 Fix: Directly use ESRGAN weight without requiring 'realesrgan'
-from model.esrgan import RealESRGANer
-upsampler = RealESRGANer(scale=4, model_path=MODEL_PATH, model=None, tile=0, tile_pad=10, pre_pad=0,
-                         half=True if torch.cuda.is_available() else False)
+# Initialize the RealESRGANer instance
+realesrgan = RealESRGANer(
+    scale=4,  # Assuming a scale of 4x
+    model_path=MODEL_PATH,
+    device='cuda' if torch.cuda.is_available() else 'cpu',  # Use GPU if available
+    half=False  # Set to True if you want to use half precision
+)
 
 def upscale_image(input_path, output_path):
     """Upscales an image using Real-ESRGAN."""
@@ -43,7 +36,7 @@ def upscale_image(input_path, output_path):
 
         # Apply super-resolution
         print(f"[INFO] Upscaling image: {input_path}")
-        output, _ = upsampler.enhance(np.array(img_pil), outscale=2)
+        output, _ = realesrgan.enhance(np.array(img_pil), outscale=2)
         output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
 
         # Save upscaled image
@@ -59,4 +52,4 @@ if __name__ == "__main__":
     parser.add_argument("-o", "--output", type=str, required=True, help="Path to save the upscaled image")
     args = parser.parse_args()
 
-    upscale_image(args.input, args.output)
+    upscale_image(args.input, args.output)