import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
import gradio as gr
from torchvision.utils import save_image

# ================================================
# FAST NEURAL STYLE MODEL (AdaIN) - Pytorch Hub
# ================================================

# Tutaj pobieramy model z PyTorch Hub (dynamic style transfer)
# Uwaga: model pobiera content + style image
model = torch.hub.load('pytorch/examples', 'fast_neural_style', source='github', model='candy')  # przykładowy styl
model.eval()

# ================================================
# UTILS
# ================================================

def preprocess(img, size=512):
    """Konwersja PIL -> Tensor"""
    transform = transforms.Compose([
        transforms.Resize(size),
        transforms.ToTensor(),
        transforms.Lambda(lambda x: x.mul(255))
    ])
    img_t = transform(img).unsqueeze(0)  # dodaj batch dim
    return img_t

def postprocess(tensor):
    """Tensor -> PIL Image"""
    tensor = tensor.clamp(0, 255).squeeze(0)
    tensor = tensor / 255
    return transforms.ToPILImage()(tensor)

# ================================================
# INFERENCE
# ================================================

def infer(content_img, style_img):
    # Preprocess
    content = preprocess(content_img)
    style = preprocess(style_img)
    
    # Użyj modelu dynamic style transfer
    with torch.no_grad():
        output = model(content, style)  # content + style
    
    return postprocess(output)

# ================================================
# GRADIO INTERFACE
# ================================================

title = "🎨 Dynamic Neural Style Transfer"
description = "Upload a content image and a style image to apply style transfer dynamically."

app = gr.Interface(
    fn=infer,
    inputs=[
        gr.Image(type="pil", label="Content Image"),
        gr.Image(type="pil", label="Style Image")
    ],
    outputs=gr.Image(label="Stylized Image"),
    title=title,
    description=description,
    allow_flagging="never"
)

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860, enable_queue=True)