Update app.py

app.py

@@ -14,15 +14,13 @@ 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])  
+# 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")
-# sys.path.append(model_folder)
-# sys.path.insert(0, model_folder)
-# print(sys.path)
 
 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()
@@ -33,7 +31,6 @@ if torch.cuda.is_available():
 else:
     torch_dtype = torch.float32
 
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
@@ -45,6 +42,7 @@ def infer(
 ):
     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])
@@ -52,119 +50,148 @@ def infer(
 
     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 = 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
 
 
-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_ = r"""
-<h1 style="text-align: center;"><b>UltraFusion</b></h1>
-
-- 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 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.
-
-- The two input images should have the same resolution; otherwise, an error will be reported.
-
-- 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(_HEADER_)
-        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, 
+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 📸✨</b></h1>""")
+            gr.Markdown("""<h1 style="text-align: center; font-size: 24px;"><b>How to Capture Short and Long Exposure Images</b></h1>""")
+            with gr.Row():
+                gr.Image("ui/en-short.png", width=IMG_W*2//3, show_label=False, interactive=False, show_download_button=False) #, height=IMG_H*2
+                gr.Image("ui/en-tap.png", width=IMG_W*2//3, show_label=False, interactive=False, show_download_button=False)
+                gr.Image("ui/en-long.png", width=IMG_W*2//3, show_label=False, interactive=False, show_download_button=False)
+            
+            with gr.Row():
+                gr.Markdown("""<h1 style="text-align: center; font-size: 12px;"><b>➁ 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 to reveal focus and exposure adjustment buttons ☀︎.</b></h1>""")
+                gr.Markdown("""<h1 style="text-align: center; font-size: 12px;"><b>➂ Drag the ☀︎ icon upward to capture a photo with a longer exposure.</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
+                )
         
-        result = gr.Image(label="Result", show_label=True, 
-                          type='pil', 
-                          image_mode='RGB', 
-                          format="png",
-                          width=IMG_W*2,
-                          height=IMG_H*2,
+            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,],
                         )
-        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,],
-                    )
+    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)