# -*- 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")
# 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, local_dir="/home/user/app")
# sys.path.append(model_folder)
# sys.path.insert(0, model_folder)
# print(sys.path)

from ultrafusion_utils import load_model, run_ultrafusion, check_input

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=60) #[uncomment to use ZeroGPU]
def infer(
    under_expo_img,
    over_expo_img,
    num_inference_steps
):
    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])
    under_expo_img, over_expo_img = check_input(under_expo_img, over_expo_img, max_l=1500)

    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, steps=num_inference_steps, 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.Accordion("Advanced Settings", open=True):
            num_inference_steps = gr.Slider(
                label="Number of inference steps",
                minimum=2,
                maximum=50,
                step=1,
                value=20,  # Replace with defaults that work for your model
                interactive=True
            )
       
        gr.Examples(
            examples=examples, 
            inputs=[under_expo_img, over_expo_img, num_inference_steps], 
            label="Examples",
            # examples_per_page=10,
            fn=infer,
            cache_examples=True,
            outputs=[result,],
            )
        # gr.Markdown(_CITE_)
    run_button.click(fn=infer,
                     inputs=[under_expo_img, over_expo_img, num_inference_steps],
                     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)