Update app.py

app.py

@@ -2,26 +2,41 @@ import gradio as gr
 from rembg import remove
 from PIL import Image
 import io
+import torch
+import numpy as np
+from RealESRGAN import RealESRGAN
 
-def background_remover(input_image):
-    output_img = remove(input_image)  # Remove the background using the rembg library
-    # No need to convert to BytesIO, just return the PIL Image object
-    return output_img
+# Initialize RealESRGAN model
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = RealESRGAN(device, scale=4)
+model.load_weights('weights/RealESRGAN_x4.pth', download=True)
+
+def background_remover_and_upscale(input_image):
+    # Remove background
+    output_img = remove(input_image)
+    
+    # Convert back to PIL Image for upscaling
+    output_img_pil = Image.fromarray(output_img)
+    
+    # Upscale image
+    sr_image = model.predict(output_img_pil)
+    
+    # Return the upscaled image directly
+    return sr_image
 
 description = """
-This app uses the rembg library to remove the background from uploaded images.
-Simply upload your image and let the model do its work. Download the result immediately after!
+This app uses the rembg library to remove the background from uploaded images and then the RealESRGAN model to upscale the image.
+Simply upload your image and let the models do their work. Download the result immediately after!
 """
 
 iface = gr.Interface(
-    fn=background_remover,
+    fn=background_remover_and_upscale,
     inputs=gr.Image(type='pil', label="Upload Image"),
     outputs="image",
     examples=["woman.jpg", "groot.jpg"],
-    title="Background Remover",
+    title="Background Remover and Image Upscaler",
     description=description,
     theme="soft"
 )
 
-# Launch the interface
-iface.launch(share=True)
+iface.launch(share=True)