support docker

app_docker.py

@@ -0,0 +1,207 @@
+# -*- 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")
+
+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)