import os
from pathlib import Path

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


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)

# Initialize the upscaler
upscaler = UpscalerESRGAN(
    model_path=Path(
        hf_hub_download(
            repo_id="philz1337x/upscaler",
            filename="4x-UltraSharp.pth",
            # revision="011deacac8270114eb7d2eeff4fe6fa9a837be70",
        )
    ),
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
    dtype=torch.float32,
)

torch.cuda.empty_cache()


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

    # Set seed for reproducibility
    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 (only if guidance is enabled)
    uncond_emb = torch.zeros_like(cond_emb) if guidance_scale > 1 else None

    # Diffusion denoising loop with mixed precision for efficiency
    with torch.autocast(device):
        for t in scheduler.timesteps:
            if guidance_scale > 1:
                # Classifier-Free Guidance (CFG)
                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:
                # Standard prediction
                noise_pred = net(x, t, cond_emb)

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

    # Decode predictions from latent space
    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)
    output_image = transforms.ToPILImage()(grid)

    # Optionally upscale the output image
    if is_upscale:
        output_image = upscaler(output_image)

    return output_image


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", height=1024, width=1024),
        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", info="Upscale output by 4x (2048x2048) using an off-the-shelf model."
        ),
    ],
    outputs=gr.Image(type="pil", label="Generated Garment", height=1024, width=1024),
    title=title,
    description=description,
    article=article,
    examples=examples,
    examples_per_page=4,
    submit_btn="Generate",
)

demo.launch()