import os
import gradio as gr
import spaces
import torch
import torch.nn as nn
from diffusers import EulerDiscreteScheduler, AutoencoderKL, UNet2DConditionModel
from huggingface_hub import hf_hub_download
from transformers import SiglipImageProcessor, SiglipVisionModel
from torchvision.io import read_image
import torchvision.transforms.v2 as transforms
from torchvision.utils import make_grid


class TryOffDiff(nn.Module):
    def __init__(self):
        super().__init__()
        self.unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")
        self.transformer = torch.nn.TransformerEncoderLayer(d_model=768, nhead=8, batch_first=True)
        self.proj = nn.Linear(1024, 77)
        self.norm = nn.LayerNorm(768)

    def adapt_embeddings(self, x):
        x = self.transformer(x)
        x = self.proj(x.permute(0, 2, 1)).permute(0, 2, 1)
        return self.norm(x)

    def forward(self, noisy_latents, t, cond_emb):
        cond_emb = self.adapt_embeddings(cond_emb)
        return self.unet(noisy_latents, t, encoder_hidden_states=cond_emb).sample


class PadToSquare:
    def __call__(self, img):
        _, h, w = img.shape  # Get the original dimensions
        max_side = max(h, w)
        pad_h = (max_side - h) // 2
        pad_w = (max_side - w) // 2
        padding = (pad_w, pad_h, max_side - w - pad_w, max_side - h - pad_h)
        return transforms.functional.pad(img, padding, padding_mode="edge")


# Set device
device = "cuda" if torch.cuda.is_available() else "cpu"

# Initialize Image Encoder
img_processor = SiglipImageProcessor.from_pretrained(
    "google/siglip-base-patch16-512",
    do_resize=False,
    do_rescale=False,
    do_normalize=False
)
img_enc = SiglipVisionModel.from_pretrained("google/siglip-base-patch16-512").eval().to(device)
img_enc_transform = transforms.Compose([
    PadToSquare(),  # Custom transform to pad the image to a square
    transforms.Resize((512, 512)),
    transforms.ToDtype(torch.float32, scale=True),
    transforms.Normalize(mean=[0.5], std=[0.5]),
])

# Load TryOffDiff Model
path_model = hf_hub_download(
    repo_id="rizavelioglu/tryoffdiff",
    filename="tryoffdiff.pth",  # or one of ["ldm-1", "ldm-2", "ldm-3", ...],
    force_download=False
)
path_scheduler = hf_hub_download(
    repo_id="rizavelioglu/tryoffdiff",
    filename="scheduler/scheduler_config.json",
    force_download=False
)
net = TryOffDiff()
net.load_state_dict(torch.load(path_model, weights_only=False))
net.eval().to(device)

# Initialize VAE (only Decoder will be used)
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae").eval().to(device)
torch.cuda.empty_cache()


# Define image generation function
@spaces.GPU(duration=10)
@torch.no_grad()
def generate_image(input_image, seed=42, guidance_scale=2.0, num_inference_steps=50, is_upscale=False):
    # Configure scheduler
    scheduler = EulerDiscreteScheduler.from_pretrained(path_scheduler)
    scheduler.is_scale_input_called = True  # suppress warning
    scheduler.set_timesteps(num_inference_steps)

    # Set random seed
    generator = torch.Generator(device=device).manual_seed(seed)
    x = torch.randn(1, 4, 64, 64, generator=generator, device=device)

    # Process input image
    cond_image = img_enc_transform(read_image(input_image))
    inputs = {k: v.to(img_enc.device) for k, v in img_processor(images=cond_image, return_tensors="pt").items()}
    cond_emb = img_enc(**inputs).last_hidden_state.to(device)

    # Prepare unconditioned embeddings
    uncond_emb = torch.zeros_like(cond_emb) if guidance_scale > 1 else None

    # Denoising loop with mixed precision
    with torch.autocast(device):
        for t in scheduler.timesteps:
            if guidance_scale > 1:
                noise_pred = net(
                    torch.cat([x] * 2), t, torch.cat([uncond_emb, cond_emb])
                ).chunk(2)
                noise_pred = noise_pred[0] + guidance_scale * (noise_pred[1] - noise_pred[0])
            else:
                noise_pred = net(x, t, cond_emb)

            scheduler_output = scheduler.step(noise_pred, t, x)
            x = scheduler_output.prev_sample

    # Decode preds
    decoded = vae.decode(1 / 0.18215 * scheduler_output.pred_original_sample).sample
    images = (decoded / 2 + 0.5).cpu()             
    
    # Create grid
    grid = make_grid(images, nrow=len(images), normalize=True, scale_each=True)
    if is_upscale:
        pass
    else:
        return transforms.ToPILImage()(grid)


title = "Virtual Try-Off Generator"
description = r"""
This is the demo of the paper <a href="https://arxiv.org/abs/2411.18350">TryOffDiff: Virtual-Try-Off via High-Fidelity Garment Reconstruction using Diffusion Models</a>. 
<br>Upload an image of a clothed individual to generate a standardized garment image using TryOffDiff. 
<br> Check out the <a href="https://rizavelioglu.github.io/tryoffdiff/">project page</a> for more information.
"""
article = r"""
Example images are sampled from the `VITON-HD-test` set, which the models did not see during training.

<br>**Citation** <br>If you find our work useful in your research, please consider giving a star ⭐ and 
a citation: 
```
@article{velioglu2024tryoffdiff,
  title     = {TryOffDiff: Virtual-Try-Off via High-Fidelity Garment Reconstruction using Diffusion Models},
  author    = {Velioglu, Riza and Bevandic, Petra and Chan, Robin and Hammer, Barbara},
  journal   = {arXiv},
  year      = {2024},
  note      = {\url{https://doi.org/nt3n}}
}
```
"""
examples = [[f"examples/{img_filename}", 42, 2.0, 20, False] for img_filename in sorted(os.listdir("examples/"))]

# Create Gradio App
demo = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Image(type="filepath", label="Reference Image"),
        gr.Slider(value=42, minimum=0, maximum=1e6, step=1, label="Seed"),
        gr.Slider(value=2.0, minimum=1, maximum=5, step=0.5, label="Guidance Scale(s)", info="No guidance applied at s=1, hence faster inference."),
        gr.Slider(value=20, minimum=0, maximum=1000, step=10, label="# of Inference Steps"),
        gr.Checkbox(value=False, label="Upscale Output")
    ],
    outputs=gr.Image(type="pil", label="Generated Garment", height=512, width=512),
    title=title,
    description=description,
    article=article,
    examples=examples,
    examples_per_page=4,
)

demo.launch()