import gradio as gr
import os
import uuid
from PIL import Image
from super_image import EdsrModel, ImageLoader
from torchvision.transforms.functional import to_pil_image

# Load pretrained model from Hugging Face
model = EdsrModel.from_pretrained('eugenesiow/edsr-base', scale=2)

# Ensure output directory exists
os.makedirs("upscaled_outputs", exist_ok=True)

def upscale_ai(image: Image.Image, resolution: str):
    # Convert to Super Image format
    scale = 2 if resolution == "2x" else 4
    model.config.scale = scale
    inputs = ImageLoader.load_image(image).unsqueeze(0)  # (1, 3, H, W)

    with torch.no_grad():
        preds = model(inputs)

    upscaled = to_pil_image(preds.squeeze(0).clamp(0, 1))
    filename = f"upscaled_{uuid.uuid4().hex[:6]}.png"
    path = os.path.join("upscaled_outputs", filename)
    upscaled.save(path)
    return image, upscaled  # Return both original and result

def process_images(images, resolution):
    results = []
    for img in images:
        image = Image.open(img).convert("RGB")
        before, after = upscale_ai(image, resolution)
        results.append([before, after])  # Side-by-side
    return results

# ------------------- UI Code with Styling -------------------

custom_theme = gr.themes.Base(
    primary_hue="blue",
    secondary_hue="fuchsia",
    neutral_hue="gray",
    spacing_size="sm",
    font=["ui-sans-serif", "Segoe UI", "sans-serif"]
).set(
    body_background_fill="#f3f4f6",
    body_text_color="#1f2937",
    button_primary_background_fill="#4f46e5",
    button_primary_text_color="white",
    button_primary_background_fill_hover="#4338ca",
    input_background_fill="#ffffff",
    input_border_color="#d1d5db",
    block_title_text_color="#1f2937"
)

with gr.Blocks(title="Ultra Pixel AI Image Upscaler", theme=custom_theme, css=".gr-button {font-size: 16px;}") as app:
    gr.Markdown(
        """
        <div style='text-align: center; padding: 1.5rem; color:#1f2937;'>
            <h1 style="font-size: 2.2rem; margin-bottom: 0.2em;">🧠 Ultra Pixel AI Image Upscaler</h1>
            <p style="font-size: 1.1rem;">
                Powerful deep learning model with side-by-side comparison.<br>
                Works beautifully across mobile, desktop, and tablets ✨
            </p>
        </div>
        """
    )

    with gr.Row():
        file_input = gr.File(label="📁 Upload Images", file_types=["image"], file_count="multiple", interactive=True)
        resolution = gr.Dropdown(["2x", "4x"], value="2x", label="Choose Resolution", interactive=True)

    upscale_btn = gr.Button("🔥 Start AI Upscaling")
    output_comparison = gr.Gallery(label="🔍 Before vs After", columns=2, show_label=False)

    upscale_btn.click(fn=process_images, inputs=[file_input, resolution], outputs=output_comparison)

    gr.Markdown(
        """
        <div style='text-align:center; margin-top: 2rem; font-size: 0.9rem; color: #6b7280;'>
            <p>Enhanced with ❤️ by Ultra Pixel | Powered by <code>super-image</code></p>
        </div>
        """
    )

app.launch()