# -*- coding: utf-8 -*-
import os
import sys
import datetime
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)
model_folder = snapshot_download(repo_id=model_name, token=auth_token, local_dir="/home/user/app")
from ultrafusion_utils import load_model, run_ultrafusion, check_input
RUN_TIMES = 0
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")
# print(under_expo_img.orig_name, over_expo_img.orig_name)
# 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")
print("num_inference_steps:", num_inference_steps)
try:
if num_inference_steps is None:
num_inference_steps = 20
num_inference_steps = int(num_inference_steps)
except Exception as e:
num_inference_steps = 20
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)
global RUN_TIMES
RUN_TIMES = RUN_TIMES + 1
print("---------------------------- Using Times---------------------------------------")
print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}: Using times: {RUN_TIMES}")
return out_pil
def build_demo():
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;
_README_ = r"""
- This is an HDR algorithm that fuses two images with different exposures.
- This can fuse two images with a very large exposure difference, even up to 9 stops.
- The two input images should have the same resolution; otherwise, an error will be reported.
- We are committed to not storing any data you upload or the results of its processing.
"""
# - The maximum resolution we support is 1500 x 1500. If the images you upload are larger than this, they will be downscaled while maintaining the original aspect ratio.
# - This is only for internal testing. Do not share it publicly.
_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 xxx@gmail.com.
"""
with gr.Blocks(css=css) as demo:
with gr.Column(elem_id="col-container"):
gr.Markdown("""UltraFusion for HDR πΈβ¨
""")
# gr.Markdown("""OpenImagingLab
""")
gr.Markdown("""How do I use it?
""")
with gr.Row():
gr.Image("ui/en-short.png", width=IMG_W//3, show_label=False, interactive=False, show_download_button=False)
gr.Image("ui/en-long.png", width=IMG_W//3, show_label=False, interactive=False, show_download_button=False)
gr.Image("ui/en-run.png", width=IMG_W//3, show_label=False, interactive=False, show_download_button=False)
with gr.Row():
gr.Markdown("""β Tap the center of the camera screen, then drag the βοΈ icon downward to capture a photo with a shorter exposure.
""")
gr.Markdown("""β Tap the center of the camera screen, then drag the βοΈ icon upward to capture a photo with a longer exposure.
""")
gr.Markdown("""β Upload the short and long exposure images, then click the 'Run' button to receive the result.
""")
gr.Markdown("""Enjoy it!
""")
with gr.Row():
under_expo_img = gr.Image(label="Short Exposure Image", show_label=True,
image_mode="RGB",
sources=["upload", ],
width=IMG_W,
height=IMG_H,
type="pil"
)
over_expo_img = gr.Image(label="Long Exposure Image", 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,
)
gr.Markdown(r"""Like it? Click the button π₯ on the image to download.
""") # width="100" height="100" 
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(_README_)
# gr.Markdown(_CITE_)
run_button.click(fn=infer,
inputs=[under_expo_img, over_expo_img, num_inference_steps],
outputs=[result,],
)
return demo
if __name__ == "__main__":
demo = build_demo()
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)