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

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

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"  # Adjust as necessary
tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')
text_encoder = TextEncoder()
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()

def generate_image(enc_text):
    noise = torch.randn((1, 100, 1, 1), device="cpu")
    with torch.no_grad():
        generated_image = model(noise, enc_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

def encode_text(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    encoded_text = text_encoder(**inputs)
    return encoded_text

def on_generate_button_click(text_input):
    if text_input:
        encoded_text = encode_text(text_input)
        generated_image = generate_image(encoded_text)
        return generated_image
    return None

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Flower Image Generator")
    text_input = gr.Textbox(label="Enter a flower-related description", value="A beautiful red rose")
    generate_button = gr.Button("Generate Image")
    output_image = gr.Image(type="numpy")  # Ensure output type is correct

    generate_button.click(on_generate_button_click, inputs=text_input, outputs=output_image)

# Launch the Gradio app
demo.launch()