darkmedia-x-api / backend /api /remote_upscale.py
cybermedia's picture
Upload folder using huggingface_hub
343eed9 verified
import os
import sys
import glob
from PIL import Image
import torch
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
def setup_upscaler():
"""Configure l'upscaler Real-ESRGAN."""
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
model_path = 'RealESRGAN_x4plus.pth'
# Téléchargement automatique du modèle si absent (géré par RealESRGANer ou wget)
if not os.path.exists(model_path):
import subprocess
print("📥 Téléchargement du modèle Real-ESRGAN...")
subprocess.run(["wget", "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"])
upsampler = RealESRGANer(
scale=4,
model_path=model_path,
model=model,
tile=0,
tile_pad=10,
pre_pad=0,
half=True # Utilise FP16 pour les GPUs modernes (A10, A100, etc.)
)
return upsampler
def main():
input_dir = "input_images"
output_dir = "output_images"
os.makedirs(output_dir, exist_ok=True)
print("🚀 Initialisation de l'upscaler AI...")
upsampler = setup_upscaler()
images = glob.glob(os.path.join(input_dir, "*"))
print(f"📂 {len(images)} images à traiter.")
for img_path in images:
name = os.path.basename(img_path)
print(f"🪄 Processing {name}...")
try:
import cv2
img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
output, _ = upsampler.enhance(img, outscale=4)
cv2.imwrite(os.path.join(output_dir, name), output)
print(f"✅ Terminé : {name}")
except Exception as e:
print(f"❌ Erreur sur {name} : {e}")
if __name__ == "__main__":
# Installation des dépendances si nécessaire (sur l'instance Lambda)
import subprocess
print("📦 Installation des dépendances AI...")
subprocess.run(["pip", "install", "basicsr", "realesrgan", "opencv-python"])
main()