Update app.py

app.py

@@ -1,63 +1,36 @@
 import torch
 from PIL import Image
+import numpy as np
 from RealESRGAN import RealESRGAN
 import gradio as gr
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model_scales = {'2x': 2, '4x': 4, '8x': 8}
 
-# Load RealESRGAN models for different scales
-models = {scale: RealESRGAN(device, scale=scale) for scale in model_scales.values()}
+model = None
+scale = 4
 
-def inference(images, scale):
-    results = []
+def upscale_image(image, scale):
+    global device
+    global model
     
-    if images is None:
-        raise gr.Error("No image uploaded. Please upload at least one image.")
+    if model is None or model.scale != scale:
+        model = RealESRGAN(device, scale=scale)
+        model.load_weights('weights/RealESRGAN_x{}.pth'.format(scale), download=True)
     
-    for image in images:
-        width, height = image.size
-        if width >= 5000 or height >= 5000:
-            raise gr.Error("The image is too large.")
-    
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        
-        # Select the appropriate model based on the chosen scale
-        model = models[model_scales[scale]]
-        result = model.predict(image.convert('RGB'))
-        print(f"Image size ({device}): {scale} ... OK")
-        results.append(result)
-    
-    return results
+    sr_image = model.predict(image)
+    return sr_image
+
+inputs = [
+    gr.Image(label="Input Image", type="pil"),
+    gr.Slider(minimum=2, maximum=8, value=4, label="Upscale Scale", step=2)
+]
+
+output = gr.Image(label="Upscaled Image")
 
-title = "✨ Real ESRGAN UpScale: 2x 4x 8x"
-description = (
-    "This advanced demo for Real-ESRGAN allows you to upscale multiple images with different models and resolutions. \nChoose the scale and upload images for high-resolution enhancement."
-)
-article = (
-    "<div style='text-align: center;'>Youtube "
-    "<a href='https://youtube.com/@kvikontent' target='_blank'>KVI Kontent</a> | "
-    "<a href='https://huggingface.co/sberbank-ai/Real-ESRGAN' target='_blank'>Model card</a></div>"
-)
+examples = [
+    ['groot.jpg', '4'],
+    ['woman.jpg', '8']
+]
 
-gr.Interface(
-    inference,
-    [
-        gr.Image(type="pil", label="Upload Image"),
-        gr.Radio(
-            list(model_scales.keys()),
-            type="value",
-            value='2x',
-            label='Resolution model',
-        ),
-    ],
-    gr.Image(type="pil", label="Output"),
-    title=title,
-    description=description,
-    article=article,
-    examples=[['groot.jpg', '2x'], ['woman.jpg', '4x']],
-    allow_flagging='never',
-    cache_examples=False,
-    theme="soft"
-).launch()
+gr.Interface(fn=upscale_image, inputs=inputs, outputs=output, title="Image Upscaler",
+             examples=examples, theme="soft", server_name="0.0.0.0").launch()