import torch
from PIL import Image
import gradio as gr
import os
from torchvision import transforms

# Download the BeautyGAN repository directly from GitHub
!pip install git+https://github.com/Honlan/BeautyGAN.git

# Import the necessary model and utils
from model import BeautyGAN

# Initialize and load the pre-trained model
model = BeautyGAN()
model.load_state_dict(torch.load("models/beautygan.pth"))
model.eval()

# Define image transformation (to match input requirements of the model)
transform = transforms.Compose([
    transforms.Resize((256, 256)),  # Resize image to fit the model's expected input size
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),  # Normalize to [-1, 1]
])

# Inference function for BeautyGAN
def beautygan_inference(image):
    # Apply the transformation to the image
    image_tensor = transform(image).unsqueeze(0)  # Add batch dimension
    
    # Run the image through the BeautyGAN model
    with torch.no_grad():
        output_image_tensor = model(image_tensor)  # Get the output from the model
    
    # Convert the output tensor back to a PIL image
    output_image = output_image_tensor.squeeze(0).clamp(0, 1).permute(1, 2, 0).numpy()
    output_image = (output_image * 255).astype("uint8")
    output_image = Image.fromarray(output_image)
    
    return output_image

# Create the Gradio interface
iface = gr.Interface(
    fn=beautygan_inference,
    inputs=gr.Image(type="pil"),  # Image input in PIL format
    outputs=gr.Image(type="pil"),  # Image output in PIL format
    live=True  # Enable live updates
)

# Launch the Gradio interface
iface.launch()