import os
os.system("pip install --force-reinstall numpy")

from rembg import remove
from realesrgan import RealESRGAN
import torch
import requests
import numpy as np
from PIL import Image
import gradio as gr

from rembg import remove
from realesrgan import RealESRGAN

# Model URL and Path
MODEL_URL = "https://huggingface.co/lllyasviel/Annotators/resolve/main/RealESRGAN_x4plus.pth"
MODEL_PATH = "RealESRGAN_x4plus.pth"

# Download model weights if not present
def download_model():
    if not os.path.exists(MODEL_PATH):
        print("Downloading Real-ESRGAN model...")
        response = requests.get(MODEL_URL)
        with open(MODEL_PATH, "wb") as f:
            f.write(response.content)

download_model()

# Load ESRGAN model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = RealESRGAN(device, scale=4)
model.load_weights(MODEL_PATH)

# ----- Feature 1: Background Removal -----
def remove_bg(input_img):
    pil_img = Image.fromarray(input_img)
    result = remove(pil_img)
    return result

# ----- Feature 2: Image Upscaling to 8K -----
def upscale_image(input_img):
    pil_img = Image.fromarray(input_img).convert("RGB")
    upscaled = model.predict(pil_img)
    upscaled = upscaled.resize((7680, 4320), Image.LANCZOS)
    return np.array(upscaled)

# UI Tabs
with gr.Blocks() as demo:
    gr.Markdown("## 🖼️ Image Tools: Background Removal & 8K Upscaling")

    with gr.Tabs():
        with gr.Tab("🧼 Remove Background"):
            with gr.Row():
                input_rm = gr.Image(type="numpy", label="Upload Image")
                output_rm = gr.Image(type="pil", label="Image w/o Background")
            btn_rm = gr.Button("Remove Background")
            btn_rm.click(remove_bg, inputs=input_rm, outputs=output_rm)

        with gr.Tab("🚀 Upscale to 8K"):
            with gr.Row():
                input_up = gr.Image(type="numpy", label="Upload Image")
                output_up = gr.Image(type="numpy", label="8K Image")
            btn_up = gr.Button("Upscale")
            btn_up.click(upscale_image, inputs=input_up, outputs=output_up)

if __name__ == "__main__":
    demo.launch()