app_docker.py

# -*- 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 -------------------------------
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)

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()

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>[email protected]</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)