Update app.py

app.py

@@ -1,35 +1,28 @@
 import gradio as gr
 from PIL import Image
-import numpy as np
-from realesrgan import RealESRGANer
-import os
+import torch
+from torchvision import transforms
 
-# Проверка наличия модели
-def check_model_path(model_path):
-    if not os.path.exists(model_path):
-        raise FileNotFoundError(f"Модель не найдена по пути: {model_path}. Скачайте модель и поместите её в папку 'weights'.")
-
-# Инициализация модели Real-ESRGAN
+# Загрузка предобученной модели ESRGAN через torch.hub
 def load_model():
-    model_path = "weights/RealESRGAN_x4plus.pth"  # Путь к модели
-    check_model_path(model_path)  # Проверка наличия модели
-    
-    model = RealESRGANer(
-        scale=4,
-        model_path=model_path,
-        model=None,
-        tile=0,
-        tile_pad=10,
-        pre_pad=0,
-        half=False
-    )
+    model = torch.hub.load('pytorch/vision:v0.10.0', 'esrgan', pretrained=True)
+    model.eval()
     return model
 
 # Функция для увеличения разрешения изображения
 def upscale_image(image):
-    img = np.array(image)
-    output, _ = model.enhance(img, outscale=4)
-    return Image.fromarray(output)
+    # Преобразование изображения в тензор
+    input_image = transforms.ToTensor()(image).unsqueeze(0)
+    
+    # Увеличение разрешения с помощью модели
+    with torch.no_grad():
+        output_image = model(input_image)
+    
+    # Преобразование тензора обратно в изображение
+    output_image = output_image.squeeze(0).clamp(0, 1)
+    output_image = transforms.ToPILImage()(output_image)
+    
+    return output_image
 
 # Загрузка модели
 model = load_model()
@@ -39,8 +32,8 @@ iface = gr.Interface(
     fn=upscale_image,
     inputs=gr.Image(type="pil", label="Входное изображение"),
     outputs=gr.Image(type="pil", label="Увеличенное изображение"),
-    title="AI Image Upscaler",
-    description="Загрузите изображение, чтобы увеличить его разрешение с помощью Real-ESRGAN."
+    title="AI Image Upscaler (ESRGAN)",
+    description="Загрузите изображение, чтобы увеличить его разрешение с помощью ESRGAN."
 )
 
 # Запуск интерфейса