import gradio as gr
from transformers import XLNetTokenizer, XLNetModel
import torch
import torch.nn as nn
import numpy as np

# TextEncoder class
class TextEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.transformer = XLNetModel.from_pretrained("xlnet-base-cased")

    def forward(self, input_ids, token_type_ids, attention_mask):
        hidden = self.transformer(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask).last_hidden_state
        context = hidden.mean(dim=1)
        context = context.view(*context.shape, 1, 1)
        return context

# Generator class
class Generator(nn.Module):
    def __init__(self, nz=100, ngf=64, nt=768, nc=3):
        super().__init__()
        self.layer1 = nn.Sequential(
            nn.ConvTranspose2d(nz+nt, ngf*8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf*8)
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(ngf*8, ngf*2, 1, 1),
            nn.Dropout2d(inplace=True),
            nn.BatchNorm2d(ngf*2),
            nn.ReLU(True)
        )
        self.layer3 = nn.Sequential(
            nn.Conv2d(ngf*2, ngf*2, 3, 1, 1),
            nn.Dropout2d(inplace=True),
            nn.BatchNorm2d(ngf*2),
            nn.ReLU(True)
        )
        self.layer4 = nn.Sequential(
            nn.Conv2d(ngf*2, ngf*8, 3, 1, 1),
            nn.Dropout2d(inplace=True),
            nn.BatchNorm2d(ngf*8),
            nn.ReLU(True)
        )
        self.layer5 = nn.Sequential(
            nn.ConvTranspose2d(ngf*8, ngf*4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf*4),
            nn.ReLU(True)
        )
        self.layer6 = nn.Sequential(
            nn.Conv2d(ngf*4, ngf, 1, 1),
            nn.Dropout2d(inplace=True),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True)
        )
        self.layer7 = nn.Sequential(
            nn.Conv2d(ngf, ngf, 3, 1, 1),
            nn.Dropout2d(inplace=True),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True)
        )
        self.layer8 = nn.Sequential(
            nn.Conv2d(ngf, ngf*4, 3, 1, 1),
            nn.Dropout2d(inplace=True),
            nn.BatchNorm2d(ngf*4),
            nn.ReLU(True)
        )
        self.layer9 = nn.Sequential(
            nn.ConvTranspose2d(ngf*4, ngf*2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf*2),
            nn.ReLU(True)
        )
        self.layer10 = nn.Sequential(
            nn.ConvTranspose2d(ngf*2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True)
        )
        self.layer11 = nn.Sequential(
            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
            nn.Tanh()
        )

    def forward(self, noise, encoded_text):
        x = torch.cat([noise, encoded_text], dim=1)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x)
        x = self.layer6(x)
        x = self.layer7(x)
        x = self.layer8(x)
        x = self.layer9(x)
        x = self.layer10(x)
        x = self.layer11(x)
        return x


# Load the model and tokenizer
model_path = "./checkpoint.pth"
tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')
text_encoder = XLNetModel.from_pretrained('xlnet-base-cased')
model = Generator()
model_state_dict = torch.load(model_path, map_location="cpu")
generator = model_state_dict['models']['generator']
model.load_state_dict(generator)

text_encoder.to("cpu")
model.to("cpu")
model.eval()

# Functions to encode text and generate image
def encode_text(text):
    text_encoder_model = TextEncoder()
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    encoded_text = text_encoder_model(**inputs)
    return encoded_text

def generate_image(text):
    encoded_text = encode_text(text)
    noise = torch.randn((1, 100, 1, 1), device="cpu")
    with torch.no_grad():
        generated_image = model(noise, encoded_text).detach().squeeze().cpu()
    gen_image_np = generated_image.numpy()
    gen_image_np = np.transpose(gen_image_np, (1, 2, 0))  # Change from CHW to HWC
    gen_image_np = (gen_image_np - gen_image_np.min()) / (gen_image_np.max() - gen_image_np.min())  # Normalize to [0, 1]
    gen_image_np = (gen_image_np * 255).astype(np.uint8)
    return gen_image_np

# Gradio interface
inputs = gr.inputs.Textbox(label="Enter a flower-related description", default="A beautiful red rose")
outputs = gr.outputs.Image(type="numpy", label="Generated Flower Image")

gr.Interface(fn=generate_image, inputs=inputs, outputs=outputs, title="Flower Image Generator", description="Enter a description of a flower to generate an image.").launch()