Spaces:
Sleeping
Sleeping
| # -*- coding: utf-8 -*- | |
| import os | |
| import sys | |
| import datetime | |
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| import spaces #[uncomment to use ZeroGPU] | |
| import torch | |
| from torchvision.transforms import ToTensor, ToPILImage | |
| # -------------------------- HuggingFace ------------------------------- | |
| # 1. Download the model online | |
| # from huggingface_hub import hf_hub_download, snapshot_download | |
| # model_name = "iimmortall/UltraFusion" | |
| # auth_token = os.getenv("HF_AUTH_TOKEN") | |
| # model_folder = snapshot_download(repo_id=model_name, token=auth_token, local_dir="/home/user/app", force_download=True) | |
| # model_folder = "" | |
| # 2. using pre-download model | |
| # from huggingface_hub import hf_hub_download, snapshot_download | |
| # model_name = "iimmortall/UltraFusion" | |
| # auth_token = os.getenv("HF_AUTH_TOKEN") | |
| # model_folder = snapshot_download(repo_id=model_name, token=auth_token, local_dir="/data", force_download=True) | |
| model_folder = "/data" | |
| sys.path.append(f"{model_folder}") | |
| from ultrafusion_utils import load_model, run_ultrafusion, check_input | |
| PYCUDA_FLAG = True | |
| try : | |
| import pycuda | |
| except Exception: | |
| PYCUDA_FLAG = False | |
| print("No pycuda!!!") | |
| RUN_TIMES = 0 | |
| to_tensor = ToTensor() | |
| to_pil = ToPILImage() | |
| ultrafusion_pipe, flow_model = load_model(model_folder) | |
| 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_lr, over_expo_img_lr, under_expo_img, over_expo_img, use_bgu = check_input(under_expo_img, over_expo_img, max_l=1500) | |
| global PYCUDA_FLAG | |
| if not PYCUDA_FLAG and use_bgu: | |
| print("No pycuda, do not run BGU.") | |
| use_bgu = False | |
| ue = to_tensor(under_expo_img_lr).unsqueeze(dim=0).to("cuda") | |
| oe = to_tensor(over_expo_img_lr).unsqueeze(dim=0).to("cuda") | |
| ue_hr = to_tensor(under_expo_img).unsqueeze(dim=0).to("cuda") | |
| oe_hr = 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, ue_hr, oe_hr, use_bgu, 'test', flow_model=flow_model, pipe=ultrafusion_pipe, steps=num_inference_steps, consistent_start=None, test_bs=8) | |
| 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 <b>xxx@gmail.com</b>. | |
| """ | |
| with gr.Blocks(css=css) as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown("""<h1 style="text-align: center; font-size: 32px;"><b>UltraFusion HDR 📸✨</b></h1>""") | |
| # gr.Markdown("""<h1 style="text-align: center; font-size: 32px;"><b>OpenImagingLab</b></h1>""") | |
| gr.Markdown("""<h1 style="text-align: center; font-size: 24px;"><b>How do I use it?</b></h1>""") | |
| 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("""<h1 style="text-align: center; font-size: 12px;"><b>➀ Tap the center of the camera screen, then drag the ☀︎ icon downward to capture a photo with a shorter exposure.</b></h1>""") | |
| gr.Markdown("""<h1 style="text-align: center; font-size: 12px;"><b>➁ Tap the center of the camera screen, then drag the ☀︎ icon upward to capture a photo with a longer exposure.</b></h1>""") | |
| gr.Markdown("""<h1 style="text-align: center; font-size: 12px;"><b>➂ Upload the short and long exposure images, then click the 'Run' button to receive the result. </b></h1>""") | |
| gr.Markdown("""<h1 style="text-align: center; font-size: 24px;"><b>Enjoy it!</b></h1>""") | |
| 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"""<h1 style="text-align: center; font-size: 18px;"><b>Like it? Click the button 📥 on the image to download.</b></h1>""") # width="100" height="100" <img src="ui/download.svg" alt="download"> | |
| 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(server_name="0.0.0.0", server_port=7860, share=True) | |
| # demo.launch(server_name="0.0.0.0", debug=True, show_api=True, show_error=True, share=False) | |