# -*- coding: utf-8 -*-
import os
import sys
import gradio as gr
import numpy as np
import random
import spaces #[uncomment to use ZeroGPU]
# from diffusers import DiffusionPipeline
import torch
from torchvision.transforms import ToTensor, ToPILImage
import logging
logging.getLogger("huggingface_hub").setLevel(logging.CRITICAL)
from huggingface_hub import hf_hub_download, snapshot_download

model_name = "iimmortall/UltraFusion"
auth_token = os.getenv("HF_AUTH_TOKEN")
# print(auth_token)
# greet_file = hf_hub_download(repo_id=model_name, filename="main.py", use_auth_token=auth_token) 
# sys.path.append(os.path.split(greet_file)[0])  
model_folder = snapshot_download(repo_id=model_name, token=auth_token)
# sys.path.append(model_folder)
sys.path.insert(0, model_folder)
print(sys.path)
# exit()

from ultrafusion_utils import load_model, run_ultrafusion

to_tensor = ToTensor()
to_pil = ToPILImage()
ultrafusion_pipe, flow_model = load_model()

device = "cuda" if torch.cuda.is_available() else "cpu"
if torch.cuda.is_available():
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32


MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

@spaces.GPU(duration=10) #[uncomment to use ZeroGPU]
def infer(
    under_expo_img,
    over_expo_img,
    # progress=gr.Progress(track_tqdm=True),
):
    # if randomize_seed:
    #     seed = random.randint(0, MAX_SEED)
    # generator = torch.Generator().manual_seed(seed)

    # image = pipe(
    #     prompt=prompt,
    #     negative_prompt=negative_prompt,
    #     guidance_scale=guidance_scale,
    #     num_inference_steps=num_inference_steps,
    #     width=width,
    #     height=height,
    #     generator=generator,
    # ).images[0]

    print(under_expo_img.size)
    print("reciving image")

    under_expo_img = under_expo_img.resize([1500, 1000])
    over_expo_img = over_expo_img.resize([1500, 1000])

    ue = to_tensor(under_expo_img).unsqueeze(dim=0).to("cuda")
    oe = to_tensor(over_expo_img).unsqueeze(dim=0).to("cuda")

    out = run_ultrafusion(ue, oe, 'test', flow_model=flow_model, pipe=ultrafusion_pipe, consistent_start=None)

    out = out.clamp(0, 1).squeeze()
    out_pil = to_pil(out)

    return out_pil


examples= [
    [os.path.join("examples", img_name, "ue.jpg"), 
     os.path.join("examples", img_name, "oe.jpg")] for img_name in sorted(os.listdir("examples"))
]
IMG_W = 320
IMG_H = 240
css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""
# max-heigh: 1500px;

_HEADER_ = '''
<h2><b>Official 🤗 UltraHDR Demo</b></h2><h2><a href='' target='_blank'><b>UltraHDR: xxx</b></a></h2>
'''

_CITE_ = r"""
📝 **Citation**

If you find our work useful for your research or applications, please cite using this bibtex:
```bibtex
@article{xxx,
  title={xxx},
  author={xxx},
  journal={arXiv preprint arXiv:xx.xx},
  year={2024}
}
```

📋 **License**

CC BY-NC 4.0. LICENSE.

📧 **Contact**

If you have any questions, feel free to open a discussion or contact us at <b>xxx@gmail.com</b>.
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(" # UltraHDR")
        with gr.Row():
            under_expo_img = gr.Image(label="UnderExposureImage", show_label=True,
                image_mode="RGB",
                sources=["upload", ],
                width=IMG_W,
                height=IMG_H,
                type="pil"
            )
            over_expo_img = gr.Image(label="OverExposureImage", show_label=True, 
                image_mode="RGB",
                sources=["upload", ],
                width=IMG_W,
                height=IMG_H,
                type="pil"
            )
        with gr.Row():
            run_button = gr.Button("Run", variant="primary") # scale=0, 
        
        result = gr.Image(label="Result", show_label=True, 
                          type='pil', 
                          image_mode='RGB', 
                          format="png",
                          width=IMG_W*2,
                          height=IMG_H*2,
                        )

        # with gr.Row():
        # prompt = gr.Text(
        #     label="Prompt",
        #     show_label=False,
        #     max_lines=1,
        #     placeholder="Enter your prompt",
        #     container=False,
        # )
        # negative_prompt = gr.Text(
        #     label="Negative prompt",
        #     max_lines=1,
        #     placeholder="Enter a negative prompt",
        #     visible=False,
        # )
        # with gr.Accordion("Advanced Settings", open=False):
        #     negative_prompt = gr.Text(
        #         label="Negative prompt",
        #         max_lines=1,
        #         placeholder="Enter a negative prompt",
        #         visible=False,
        #     )

        #     seed = gr.Slider(
        #         label="Seed",
        #         minimum=0,
        #         maximum=MAX_SEED,
        #         step=1,
        #         value=0,
        #     )

        #     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

        #     with gr.Row():
        #         width = gr.Slider(
        #             label="Width",
        #             minimum=256,
        #             maximum=MAX_IMAGE_SIZE,
        #             step=32,
        #             value=1024,  # Replace with defaults that work for your model
        #         )

        #         height = gr.Slider(
        #             label="Height",
        #             minimum=256,
        #             maximum=MAX_IMAGE_SIZE,
        #             step=32,
        #             value=1024,  # Replace with defaults that work for your model
        #         )

        #     with gr.Row():
        #         guidance_scale = gr.Slider(
        #             label="Guidance scale",
        #             minimum=0.0,
        #             maximum=10.0,
        #             step=0.1,
        #             value=0.0,  # Replace with defaults that work for your model
        #         )

        #         num_inference_steps = gr.Slider(
        #             label="Number of inference steps",
        #             minimum=1,
        #             maximum=50,
        #             step=1,
        #             value=2,  # Replace with defaults that work for your model
        #         )
       
        gr.Examples(
            examples=examples, 
            inputs=[under_expo_img, over_expo_img], 
            label="Examples",
            # examples_per_page=10,
            cache_examples=False,
            # fn=infer,
            )
        # gr.Markdown(_CITE_)
    run_button.click(fn=infer,
                     inputs=[under_expo_img, over_expo_img],
                     outputs=[result,],
                    )

if __name__ == "__main__":
    demo.queue(max_size=10)
    demo.launch(share=True)
    # demo.launch(server_name="0.0.0.0", debug=True, show_api=True, show_error=True, share=False)