Zip files before uploading

app.py

@@ -6,6 +6,9 @@ import json
 import os
 import random
 import tempfile
+import shortuuid
+from apscheduler.schedulers.background import BackgroundScheduler
+import shutil
 
 import gradio as gr
 import torch
@@ -105,15 +108,18 @@ def upload_stage1_result(stage1_param_path: str, stage1_result_path: str,
     if not UPLOAD_REPO_ID:
         return
     try:
-        random_folder = random.randint(0,1000)
-        hf_api.upload_file(path_or_fileobj=stage1_param_path,
-                           path_in_repo=f'stage1_params/{random_folder}/{save_name}.json',
-                           repo_id=UPLOAD_REPO_ID,
-                           repo_type='dataset')
-        hf_api.upload_file(path_or_fileobj=stage1_result_path,
-                           path_in_repo=f'stage1_results/{random_folder}/{save_name}.pth',
-                           repo_id=UPLOAD_REPO_ID,
-                           repo_type='dataset')
+        folder_params = "tmp/results/stage1_params"
+        folder_results = "tmp/results/stage1_results"
+        
+        path_params = f"{folder_params}/{save_name}.json"
+        path_results = f"{folder_results}/{save_name}.pth"
+        
+        os.makedirs(folder_params, exist_ok=True)
+        os.makedirs(folder_results, exist_ok=True)
+        
+        shutil.copy(stage1_param_path, path_params)
+        shutil.copy(stage1_result_path, path_results)
+
     except Exception as e:
         print(e)
 
@@ -140,18 +146,20 @@ def upload_stage2_info(stage1_param_file_hash_name: str,
         param_file.write(json.dumps(stage2_params))
     stage2_param_file_hash_name = get_param_file_hash_name(param_file.name)
     save_name = f'{stage1_param_file_hash_name}_{stage2_param_file_hash_name}'
-
+    
     try:
-        random_folder = random.randint(0,1000)
-        hf_api.upload_file(path_or_fileobj=param_file.name,
-                           path_in_repo=f'stage2_params/{random_folder}/{save_name}.json',
-                           repo_id=UPLOAD_REPO_ID,
-                           repo_type='dataset')
+        folder_params = "tmp/results/stage2_params"
+
+        os.makedirs(folder_params, exist_ok=True)
+        path_params = f"{folder_params}/{save_name}.json"
+        shutil.copy(param_file.name, path_params)
+        
         if UPLOAD_RESULT_IMAGE:
-            hf_api.upload_file(path_or_fileobj=stage2_output_path,
-                               path_in_repo=f'stage2_results/{random_folder}/{save_name}.png',
-                               repo_id=UPLOAD_REPO_ID,
-                               repo_type='dataset')
+            folder_results = "tmp/results/stage2_results"
+            os.makedirs(folder_results, exist_ok=True)
+            path_results = f"{folder_results}/{save_name}.png"
+            shutil.copy(stage2_output_path, path_results)
+
     except Exception as e:
         print(e)
 
@@ -188,17 +196,16 @@ def upload_stage2_3_info(stage1_param_file_hash_name: str,
     save_name = f'{stage1_param_file_hash_name}_{stage2_3_param_file_hash_name}'
 
     try:
-        random_folder = random.randint(0,1000)
-        hf_api.upload_file(path_or_fileobj=param_file.name,
-                           path_in_repo=f'stage2_3_params/{random_folder}/{save_name}.json',
-                           repo_id=UPLOAD_REPO_ID,
-                           repo_type='dataset')
+        folder_params = "tmp/results/stage2_3_params"
+        os.makedirs(folder_params, exist_ok=True)
+        path_params = f"{folder_params}/{save_name}.json"
+        shutil.copy(param_file.name, path_params)
+
         if UPLOAD_RESULT_IMAGE:
-            hf_api.upload_file(
-                path_or_fileobj=stage2_3_output_path,
-                path_in_repo=f'stage2_3_results/{random_folder}/{save_name}.png',
-                repo_id=UPLOAD_REPO_ID,
-                repo_type='dataset')
+            folder_results = "tmp/results/stage2_3_results"
+            os.makedirs(folder_results, exist_ok=True)
+            path_results = f"{folder_results}/{save_name}.png"
+            shutil.copy(stage2_3_output_path, path_results)
     except Exception as e:
         print(e)
 
@@ -221,6 +228,21 @@ def show_gallery_view() -> tuple[dict, dict]:
 def show_upscaled_view() -> tuple[dict, dict]:
     return _update_result_view(False)
 
+def upload_files():
+    """Zips files and uploads to dataset. Local data is deleted
+    """
+    if os.path.exists("tmp/results") and os.path.isdir("tmp/results"):
+        try:
+            shutil.make_archive("tmp/results", 'zip', "tmp/results")
+            hf_api.upload_file(
+                path_or_fileobj="tmp/results.zip",
+                path_in_repo=f"results_{shortuuid.uuid()}.zip",
+                repo_id=UPLOAD_REPO_ID,
+                repo_type="dataset",
+            )
+            shutil.rmtree("tmp/results")
+        except Exception as e:
+            print(e)
 
 examples = [
     'high quality dslr photo, a photo product of a lemon inspired by natural and organic materials, wooden accents, intricately decorated with glowing vines of led lights, inspired by baroque luxury',
@@ -669,5 +691,10 @@ with gr.Blocks(css='style.css') as demo:
         ],
         queue=False,
     )
+    
+    if UPLOAD_REPO_ID:
+        scheduler = BackgroundScheduler()
+        scheduler.add_job(func=upload_files, trigger="interval", seconds=60*5)
+        scheduler.start()
 
 demo.queue(api_open=False, max_size=MAX_QUEUE_SIZE).launch(debug=DEBUG)

requirements.txt

@@ -12,3 +12,5 @@ sentencepiece==0.1.98
 tokenizers==0.13.3
 tqdm==4.65.0
 transformers==4.28.1
+shortuuid==1.0.11
+apscheduler==3.10.1