import os

os.system("pip freeze")
import spaces
import tempfile
import shutil

import gradio as gr
import torch as torch
from gradio_dualvision import DualVisionApp
from huggingface_hub import login
from PIL import Image
from windowseat_inference import load_network, run_inference

uri_base = "Qwen/Qwen-Image-Edit-2509"
uri_lora = "huawei-bayerlab/windowseat-reflection-removal-v1-0"

if "HF_TOKEN_LOGIN" in os.environ:
    login(token=os.environ["HF_TOKEN_LOGIN"])

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

vae, transformer, embeds_dict, processing_resolution = load_network(uri_base, uri_lora, device)

# # As of transformers==4.57.1 , xformers is not supported in QwenImageTransformer2DModel
# try:
#     transformer.enable_xformers_memory_efficient_attention()
#     print("xformers enabled")
# except:
#     print("xformers not enabled")


class WindowSeatApp(DualVisionApp):
    DEFAULT_SEED = 2025

    def make_header(self):
        gr.Markdown(
            """
            ## WindowSeat Reflection Removal
            """
        )
        with gr.Row(elem_classes="remove-elements"):
            gr.Markdown(
                f"""
                <p align="center">
                <a title="Website" href="https://hf.co/spaces/huawei-bayerlab/windowseat-reflection-removal-web" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
                    <img src="https://img.shields.io/badge/%E2%99%A5%20Project%20-Website-blue">
                </a>
                <a title="arXiv" href="https://arxiv.org/abs/2512.05000" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
                    <img src="https://img.shields.io/badge/%F0%9F%93%84%20arXiv%20-Paper-AF3436">
                </a>
                <a title="Github" href="https://github.com/huawei-bayerlab/windowseat-reflection-removal" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
                    <img src="https://img.shields.io/github/stars/huawei-bayerlab/windowseat-reflection-removal?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
                </a>
                <a title="Model weights" href="https://hf.co/huawei-bayerlab/windowseat-reflection-removal-v1-0" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
                    <img src="https://img.shields.io/badge/%F0%9F%A4%97%20WindowSeat%20Model%20-Weights-yellow" alt="imagedepth">
                </a>
                <a title="Social" href="https://twitter.com/antonobukhov1" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
                    <img src="https://shields.io/twitter/follow/:?label=Subscribe%20for%20updates!" alt="social">
                </a>
                </p>                    
                <p align="center" style="margin-top: 0px;">
                    Upload a photo or pick an example below to remove reflections, wait for the result, then explore it with the slider.
                    If a quota limit appears, duplicate the space to continue.
                </p>
                """
        )

    def build_user_components(self):
        return {}

    def process(self, image_in: Image.Image, **kwargs):
        input_temp_dir = tempfile.mkdtemp()
        output_temp_dir = tempfile.mkdtemp()

        try:
            input_image_path = os.path.join(input_temp_dir, "image.png")
            image_in.save(input_image_path)
            run_inference(
                vae, 
                transformer, 
                embeds_dict, 
                processing_resolution, 
                input_temp_dir,
                output_temp_dir, 
                use_short_edge_tile=True, 
                save_comparison=False, 
                save_alternating=False,
            )
            output_image_path = os.path.join(output_temp_dir, "image_windowseat_output.png")
            result_image = Image.open(output_image_path)
            result_image.load()

            out_modalities = {
                "Result": result_image,
            }

            out_settings = {}

            return out_modalities, out_settings

        finally:
            if os.path.exists(input_temp_dir):
                shutil.rmtree(input_temp_dir)
            if os.path.exists(output_temp_dir):
                shutil.rmtree(output_temp_dir)


with WindowSeatApp(
    title="WindowSeat Reflection Removal",
    examples_path="example_images",
    examples_per_page=12,
    right_selector_visible=False,
    advanced_settings_visible=False,
    squeeze_canvas=True,
    spaces_zero_gpu_enabled=True,
) as demo:
    demo.queue(
        api_open=False,
    ).launch(
        server_name="0.0.0.0",
        server_port=7860,
        ssr_mode=False,
    )