remove persistence

app.py

@@ -25,23 +25,8 @@ desc = """
 """
 
 
-def init_persistence(purge=False):
-    if not os.path.exists('/data'):
-        return
-    os.environ['ckpt_dir'] = "/data/Marigold_ckpt"
-    os.environ['TRANSFORMERS_CACHE'] = "/data/hfcache"
-    os.environ['HF_DATASETS_CACHE'] = "/data/hfcache"
-    os.environ['HF_HOME'] = "/data/hfcache"
-    if purge:
-        os.system("rm -rf /data/Marigold_ckpt/*")
-
-
-def download_code_weights():
+def download_code():
     os.system('git clone https://github.com/prs-eth/Marigold.git')
-    os.system('cd Marigold && bash script/download_weights.sh')
-    os.system('echo /data && ls -la /data')
-    os.system('echo /data/Marigold_ckpt && ls -la /data/Marigold_ckpt')
-    os.system('echo /data/Marigold_ckpt/Marigold_v1_merged && ls -la /data/Marigold_ckpt/Marigold_v1_merged')
 
 
 def find_first_png(directory):
@@ -61,20 +46,14 @@ def marigold_process(path_input, path_out_png=None, path_out_obj=None, path_out_
     os.makedirs(path_output_dir, exist_ok=True)
     shutil.copy(path_input, path_input_dir)
 
-    persistence_args = ""
-    if os.path.exists('/data'):
-        persistence_args = "--checkpoint /data/Marigold_ckpt/Marigold_v1_merged"
-
     os.system(
         f"cd Marigold && python3 run.py "
-        f"{persistence_args} "
         f"--input_rgb_dir \"{path_input_dir}\" "
         f"--output_dir \"{path_output_dir}\" "
         f"--n_infer 10 "
         f"--denoise_steps 10 "
     )
 
-    # depth_colored, depth_bw, depth_npy
     path_out_colored = find_first_png(path_output_dir + "/depth_colored")
     assert path_out_colored is not None, "Processing failed"
     path_out_bw = find_first_png(path_output_dir + "/depth_bw")
@@ -114,13 +93,23 @@ iface = gr.Interface(
         ),
     ],
     allow_flagging="never",
-    # examples=[
-    #     [
-    #         os.path.join(os.path.dirname(__file__), "files/test.png"),
-    #         os.path.join(os.path.dirname(__file__), "files/test.png.out.png"),
-    #         os.path.join(os.path.dirname(__file__), "files/test.png.out.2.png"),
-    #     ],
-    # ],
+    examples=[
+        [
+            os.path.join(os.path.dirname(__file__), "files/bee.jpg"),
+            os.path.join(os.path.dirname(__file__), "files/bee_vis.png"),
+            os.path.join(os.path.dirname(__file__), "files/bee_pred.png"),
+        ],
+        [
+            os.path.join(os.path.dirname(__file__), "files/cat.jpg"),
+            os.path.join(os.path.dirname(__file__), "files/cat_vis.png"),
+            os.path.join(os.path.dirname(__file__), "files/cat_pred.png"),
+        ],
+        [
+            os.path.join(os.path.dirname(__file__), "files/swings.jpg"),
+            os.path.join(os.path.dirname(__file__), "files/swings_vis.png"),
+            os.path.join(os.path.dirname(__file__), "files/swings_pred.png"),
+        ],
+    ],
     css="""
     .viewport {
         aspect-ratio: 4/3;
@@ -134,6 +123,5 @@ iface = gr.Interface(
 
 
 if __name__ == "__main__":
-    init_persistence()
-    download_code_weights()
+    download_code()
     iface.queue().launch(server_name="0.0.0.0", server_port=7860)